根据列名获取pandas数据框的子集

时间:2013-04-21 00:11:17

标签: python regex pandas

我有一个包含140个样本(列)和~27000个SNP(行)的pd.DataFrame。每个列名都有一个种群名称加上一个数字(例如'FLFL04'或'MI03'),其中有6个不同的种群和不同数量的样本。

我想根据人口名称取相应种群的子集进行进一步计算(Hardy-Weinberg精确检验);我可以用循环和正则表达式做到这一点,但希望有一个更快的解决方案。 有没有办法根据列名创建子集(而不是它们的内容)?

编辑: 我目前的做法如下:

(任何pd.DataFrame都会这样做,包含以下列:

data.columns = ['FLFL04', 'FLFL08', 'FLFL08replicate', 'FLFL10', 'FLFL13', 'FLFL14', 'FLFL15', 'FLFL15replicate', 'FLFL16', 'FLFL17', 'FLFL17replicate', 'FLFL19', 'FLFL20', 'FLFL20replicate', 'FLFL21', 'FLFL23', 'FLFL26', 'FLFL28', 'FLFL28replicate', 'FLFL29', 'FLFL29replicate', 'FLFL30', 'HSPQ01', 'HSPQ01replicate', 'HSPQ01replicate2', 'HSPQ02', 'HSPQ02replicate', 'HSPQ02replicate2', 'HSPQ03', 'HSPQ04', 'HSPQ04replicate', 'HSPQ04replicate2', 'HSPQ06', 'HSPQ07', 'HSPQ08', 'HSPQ09', 'HSPQ09replicate', 'HSPQ10', 'HSPQ10replicate', 'HSPQ11', 'HSPQ12', 'HSPQ13', 'HSPQ14', 'HSPQ15', 'HSPQ16', 'HSPQ17', 'HSPQ18', 'HSPQ19', 'HSPQ21', 'HSPQ22', 'HSPQ22replicate', 'KFO1', 'KFO2', 'KFO3', 'KFO4', 'KFO5', 'KFO8', 'MI01', 'MI02', 'MI03', 'MI03replicate', 'MI03replicate2', 'MI04', 'MI05', 'MI06', 'MI06replicate', 'MI06replicate2', 'MI08', 'MI09', 'MI09replicate', 'MI09replicate2', 'MI10', 'MI11', 'MI12', 'MI12replicate', 'MI13', 'MI13replicate', 'MI14', 'MI15', 'MI16', 'MI16replicate', 'MI17', 'MI18', 'MI19', 'MI20', 'MI21', 'SFQ01', 'SFQ02', 'SFQ03', 'SFQ03replicate', 'SFQ05', 'SFQ05replicate', 'SFQ06', 'SFQ06replicate', 'SFQ07', 'SFQ08', 'SFQ08replicate', 'SFQ09', 'SFQ09replicate', 'SFQ10', 'SFQ10replicate', 'SFQ11', 'SFQ13', 'SFQ14', 'SFQ15', 'SFQ16', 'SFQ17', 'SFQ21', 'SFQ23', 'SFQ24', 'SFQ25', 'SFQ26', 'WWA01', 'WWA01replicate', 'WWA01replicate2', 'WWA03', 'WWA03replicate', 'WWA03replicate2', 'WWA04', 'WWA05', 'WWA05replicate', 'WWA05replicate2', 'WWA07', 'WWA08', 'WWA08replicate', 'WWA09', 'WWA10', 'WWA12', 'WWA17', 'WWA17replicate', 'WWA18', 'WWA21', 'WWA23', 'WWA24', 'WWA25', 'WWA25replicate', 'WWA26', 'WWA27', 'WWA28', 'WWA30']

def get_pop_subset(pop_list, pop_name): 
    pop_result_list = []
    for i, pop in enumerate(data.columns):
    curr_pop = re.findall('([A-Z]+)', pop)[0]
    if curr_pop == pop_name:
        pop_result_list.append(pop)
    return pop_result_list

pops = ['FLFL', 'HSPQ', 'KFO', 'MI', 'SFQ', 'WWA']
subsets = []
for val in pops:
    subsets.append(get_pop_subset(data.columns, val))

for val in subsets:
    print data[val]

然后我调用其他func而不是

    print data[val]

并将每个附加到新的df。 虽然这有效,但我希望能够获得更快,更有效的解决方案

感谢, 马丁

3 个答案:

答案 0 :(得分:7)

难道你不能使用带有参数“regex”的内置DataFrame方法“filter”来实现同样的功能吗?例如,

df2 = df.filter(regex='FLFL')

返回一个新的DataFrame,其中所有列都以FLFL开头。

答案 1 :(得分:4)

好的,对于你的情况,我会使用groupby。您可以使用axis=1向其传递函数以循环列(请参阅文档中的here):

>>> df 
   FLFL04  FLFL29rep  HSPQ12  MI03repl  MI16repl  SFQ10re  WWA05r
0       0          3       6         9        12       15      18
1       1          4       7        10        13       16      19
2       2          5       8        11        14       17      20
>>> df.groupby(lambda x: re.match("[A-Z]+", x).group(), axis=1)
<pandas.core.groupby.DataFrameGroupBy object at 0x9ae660c>
>>> grouped = df.groupby(lambda x: re.match("[A-Z]+", x).group(), axis=1)

然后我们可以遍历群组:

>>> for name, group in grouped:
    print 'group name:', name
    print 'dataframe:'
    print group
...     
group name: FLFL
dataframe:
   FLFL04  FLFL29rep
0       0          3
1       1          4
2       2          5
group name: HSPQ
dataframe:
   HSPQ12
0       6
1       7
2       8
group name: MI
dataframe:
   MI03repl  MI16repl
0         9        12
1        10        13
2        11        14
group name: SFQ
dataframe:
   SFQ10re
0       15
1       16
2       17
group name: WWA
dataframe:
   WWA05r
0      18
1      19
2      20

或者把它变成字典:

>>> pprint.pprint(dict(list(grouped)))
{'FLFL':    FLFL04  FLFL29rep
0       0          3
1       1          4
2       2          5,
 'HSPQ':    HSPQ12
0       6
1       7
2       8,
 'MI':    MI03repl  MI16repl
0         9        12
1        10        13
2        11        14,
 'SFQ':    SFQ10re
0       15
1       16
2       17,
 'WWA':    WWA05r
0      18
1      19
2      20}

答案 2 :(得分:0)

使用正则表达式获取列名的任何子字符串,以“FLFL”开头,后跟两位数字,删除任何不匹配项 (NA),删除重复项并排序。

df.columns.str.extract(r'(FLFL\d\d)', expand=False).dropna().drop_duplicates().sort_values()