Numpy - 大型 csv 文件的矢量化计算

Question

我有一个 20 GB 的trades.csv文件。 它有两列（trade_time 和 price）。 csv 文件包含 6.5 亿行。

样本数据

https://gist.github.com/dsstex/bc885ed04a6de98afc7102ed08b78608

熊猫数据框

df = pd.read_csv("trades.csv", index_col=0, parse_dates=True)

我想根据百分比检查价格是上涨还是下跌。 如果价格先达到 up_value（例如 +1%），则结果为 1。如果价格先达到 down_value（例如 -0.5%），则结果为 0。我需要对所有 6.5 亿行执行此操作。

目前，数据框只有两列。 trade_time(index), price 。 我想要一个名为“结果”的新列。

import pandas as pd

df = pd.read_csv("trades.csv", index_col=0, parse_dates=True)
df["result"] = None

print(df)

up_percentage = 0.2
down_percentage = 0.1


def calc_value_from_percentage(percentage, whole):
    return (percentage / 100) * whole


def set_result(index):

    up_value = 0
    down_value = 0

    for _, current_row_price, _ in df.loc[index:].itertuples():
        if up_value == 0 or down_value == 0:

            up_delta = calc_value_from_percentage(up_percentage, current_row_price)
            down_delta = calc_value_from_percentage(down_percentage, current_row_price)

            up_value = current_row_price + up_delta
            down_value = current_row_price - down_delta

        if current_row_price > up_value:
            df.loc[index, "result"] = 1
            return

        if current_row_price < down_value:
            df.loc[index, "result"] = 0
            return


for ind, _, _ in df.itertuples():
    set_result(ind)

df.to_csv("results.csv", index=True, header=True)
print(df)

结果

https://gist.github.com/dsstex/fe3759beedbf9c46ace382a7eef3d12c

注意：由于数据不足，上述文件中的大部分底部行的“结果”值为“无”。 所以该值为空白。

---

目前，我正在使用 pandas itertuples()来处理文件。 我想要一个矢量化的解决方案，因为我有一个巨大的文件。

注意：上个月我问了这个问题。 这是一个后续问题。 它与这个答案有关。 在那个答案中，作者使用了200的固定大小up_value/down_value 。 但我追求的是基于百分比的矢量化解决方案。

任何帮助是极大的赞赏。

谢谢

Answer 1

阅读您的完整代码，我终于理解了您的算法。

对于数据帧的每个索引，您必须计算“结果”是 1 还是 0，以便：

• 1表示我在数据框中找到了另一个价格，即循环中的current_row_price ，它大于我的原始价格 - 对于当前索引，在if块中计算 - 通过up_delta值；
• 0表示我在df中找到比我的原始价格低一个down_delta值的另一个价格。

我想出了这个代码。 也许循环是可以避免的，但这应该更快。

from enum import Enum

import pandas as pd


class Result(int, Enum):
    DOWN = 0
    UP = 1
    MISSING = 2


df = pd.read_csv("trades.csv", index_col=0, parse_dates=True)
df["result"] = Result.MISSING

# constants 
up_percentage = 0.2
down_percentage = 0.1

# compute upper and lower bound for every row
df["upper_bound"] = df["price"] * (1 + up_percentage / 100)
df["lower_bound"] = df["price"] * (1 - down_percentage / 100)

# for each row get current upper and lower bounds, and check 
# in all dataframe if any row is greater/lower than these values
for i, row in df.iterrows():
    series_up: pd.Series = pd.Series(df["price"].loc[i:] > row["upper_bound"])
    series_up_index = series_up[series_up].index
    series_up_min = series_up_index.min()

    series_down: pd.Series = pd.Series(df["price"].loc[i:] < row["lower_bound"])
    series_down_index = series_down[series_down].index
    series_down_min = series_down_index.min()

    is_up_hit = bool(series_up_min) and not pd.isna(series_up_min)
    is_down_hit = bool(series_down_min) and not pd.isna(series_down_min)

    if is_up_hit and is_down_hit:
        if series_up_min < series_down_min:
            result = Result.UP
        else:
            result = Result.DOWN
    elif is_up_hit:
        result = Result.UP
    elif is_down_hit:
        result = Result.DOWN
    else:
        result = Result.MISSING

    df.loc[i, "result"] = result


# remove utility columns
df.drop(columns=["upper_bound", "lower_bound"], inplace=True)

# store result
df.to_csv("results.csv", index=True, header=True)

Answer 2

原始算法非常慢，因为它正在使用 iterrows/tuples 进行嵌套循环。

如果我理解得很好，对于每一行，您检查是否有任何后行达到“固定”百分比。 如果它up ，则标记为 1，如果down ，则标记为 0，否则不标记（ None ）

我到达了这个代码。 它不是矢量化的，但它在我的机器上运行比最初的问题和接受的解决方案快得多。

可能是因为有 650M 行，它会变得更慢。

import pandas as pd
import numpy as np

from time import time

df = pd.read_csv("trades.csv", index_col=0, parse_dates=True)
t0=time()

up_percentage = 0.2
down_percentage = 0.1

# Precalculate the percentages
df['upper'] = df['price']*(1+up_percentage/100)
df['lower'] = df['price']*(1-down_percentage/100)

pupper = np.array([np.argmax(df.price.values[n:] > up_value)   for n,up_value   in enumerate(df.upper)])-1
plower = np.array([np.argmax(df.price.values[n:] < down_value) for n,down_value in enumerate(df.lower)])-1

df["result"] = None
# These two cases occur when the index is not found, but no need to re-set to None. 
# df.loc[pupper<0,'result']=None
# df.loc[plower<0,'result']=None
# If the upper value is found and it occurs before the lower, set it to 1
df.loc[(pupper>0)&((plower<0)|(pupper<plower)),'result']=1
# If the upper value is found and it occurs before the lower, set it to 1
df.loc[(pupper<0)&(plower>0),'result']=0

print(f"{1000*(time()-t0):0.2f}ms")

基准测试：只计算执行操作的时间，而不是加载/保存 CSV。

• 原版：19s
• Crissal 的：6537 毫秒
• 这段代码：135ms

检查运行原始代码 + 提议代码作为df2的相等性并比较：

df3 = df.merge(df2, left_index=True, right_index=True).fillna('empty')
print(f"The result columns are equal: {(df3.result_x==df3.result_y).all()}")