在 Python 中优化数据帧子集操作

Optimizing a dataframe subset operation in Python

总结问题

我正在尝试优化我编写的一些代码。在当前形式下,它按预期工作,但是由于脚本需要大量循环,因此需要很长时间才能 运行.

我正在寻找加速下面描述的代码的方法。

详细说明问题

在这个名为 master 的数据框中,有 3,936,192 行。 Position 列表示基因组 window。在这个数据帧中出现了 76 次。这样 master[master['Position'] == 300] returns 一个 76 行的数据框,并且每个位置的独特外观都相似。我对数据框的每个子集都做了一些操作。

可以找到数据here

我当前的代码采用以下形式:

import pandas as pd

master = pd.read_csv(data_location)

windows = sorted(set(master['Position']))

window_factor = []

               # loop through all the windows, look at the cohort of samples, ignore anything not CNV == 2
               # if that means ignore all, then drop the window entirely
               # else record the 1/2 mean of that windows normalised coverage across all samples. 

for window in windows:
    current_window = master[master['Position'] == window]

    t = current_window[current_window['CNV'] == 2]

    if t.shape[0] == 0:
        window_factor.append('drop')

    else:
        window_factor.append(
            np.mean(current_window[current_window['CNV'] == 2]['Normalised_coverage'])/2)

但是,运行 这需要非常长的时间,而且我想不出加快速度的方法,尽管我知道肯定有一个方法。

您可以做几件事:

  • 考虑使用 np.array 而不是为 window_factor 使用 python 列表,因为 你知道数组的长度。
  • t 已经是 current_window[current_window['CNV'] == 2] 在计算 np.mean 时使用 t

你也可以用profiler看看有没有开销很大的操作,或者干脆考虑用C++重新实现代码(很简单)。

我认为 .groupby() 函数是您在这里需要的:

fac = []
for name,group in master.groupby('Position'):
    if all(group['CNV'] != 2):
        fac.append('drop')
    else:
        fac.append(np.mean(group[group['CNV'] == 2]['Normalised_coverage'])/2)


我下载了你的数据master.csv,生成的数据完全一样,运行在我的笔记本电脑上,时间从 6 分钟减少到 30 秒。 希望对你有帮助。

你的 df 没有那么大,你的代码也没有什么问题:

  • 如果你使用 np.mean 并且一个值是 np.nan 它 returns np.nan
  • 计算平均值后除以2即可。
  • 在我看来这是 groupby
  • 的完美案例
  • Return 一个字符串,而其他结果是 float 你可能会考虑使用 np.nan 改为
import pandas as pd

df = pd.read_csv("master.csv")

def fun(x):
    t = x[x["CNV"]==2]
    return t["Normalised_coverage"].mean()/2

# returns np.nan when len(t)==0
out = df.groupby('Position').apply(fun)
CPU times: user 34.7 s, sys: 72.5 ms, total: 34.8 s
Wall time: 34.7 s

或者在 groupby as

之前过滤得更快
%%time
out = df[df["CNV"]==2].groupby("Position")["Normalised_coverage"].mean()/2

CPU times: user 82.5 ms, sys: 8.03 ms, total: 90.5 ms
Wall time: 87.8 ms

更新: 在最后一种情况下,如果您确实需要跟踪 df["CNV"]!=2 所在的组,您可以使用此代码:

import numpy as np
bad = df[df["CNV"]!=2]["Position"].unique()
bad = list(set(bad)-set(out.index))

out = out.reset_index(name="value")

out1 = pd.DataFrame({"Position":bad,
                     "value":[np.nan]*len(bad)})

out = pd.concat([out,out1],
                ignore_index=True)\
        .sort_values("Position")\
        .reset_index(drop=True)

这会将 160ms 添加到您的计算中。

使用 groupby 和查询是我采用的解决方案。

import pandas as pd
import numpy as np

master = pd.read_csv("/home/sean/Desktop/master.csv", index_col=0)

windows = sorted(set(master['Position']))

g = master.groupby("Position")

master.query("Position == 24386700").shape

g = master.query("CNV == 2").groupby("Position")

p = g.Normalised_coverage.mean() / 2