如何将 ta-lib 函数包装为 Polars 表达式

Question

我正在尝试通过Polars调用一些TA-lib(https://github.com/mrjbq7/ta-lib)函数，以便通过Polars的并行计算框架可以计算出多只股票的技术指标。

这是示例代码

import talib
import polars as pl
import yfinance as yf

tesla = yf.Ticker('TSLA')
tesla_data = tesla.history(period="1Y")
tesla_data["Date"]=tesla_data.index
pl_df = pl.from_pandas(tesla_data[["Date", "Open", "High", "Low", "Close", "Volume"]])

# Method 1. Using ta-lib as a direct function call.
mv_kama = talib.KAMA(pl_df["Close"], 30)

# Method 2. Using ta-lib as Polars expression
def kama30() -> pl.Expr:
    return talib.KAMA(pl.col("Close"), 30)

pl_df2 = pl_df.select([
    pl.col("Close"),
    kama30()
])

方法 2 代码片段未能运行，错误消息为：

TypeError                                 Traceback (most recent call last)
Input In [5], in <cell line: 17>()
     14 def kama30() -> pl.Expr:
     15     return talib.KAMA(pl.col("Close"), 30)
     17 pl_df2 = pl_df.select([
     18     pl.col("Close"),
---> 19     kama30()
     20 ])

Input In [5], in kama30()
     14 def kama30() -> pl.Expr:
---> 15     return talib.KAMA(pl.col("Close"), 30)

File C:\ProgramData\Anaconda3\envs\Charm3.9\lib\site-packages\talib\__init__.py:64, in _wrapper.<locals>.wrapper(*args, **kwds)
     61     _args = args
     62     _kwds = kwds
---> 64 result = func(*_args, **_kwds)
     66 # check to see if we got a streaming result
     67 first_result = result[0] if isinstance(result, tuple) else result

TypeError: Argument 'real' has incorrect type (expected numpy.ndarray, got Expr)

如果有人能建议如何正确地完成此操作，我们将不胜感激。

谢谢！

Answer 1

无需定义Polars函数表达式，只需使用map函数即可。这是工作代码

import talib
import polars as pl
import yfinance as yf

tesla = yf.Ticker('TSLA')
tesla_data = tesla.history(period="1Y")
tesla_data["Date"]=tesla_data.index
pl_df = pl.from_pandas(tesla_data[["Date", "Open", "High", "Low", "Close", "Volume"]])


pl_df2 = pl_df.select([
    pl.col("Close"),
    pl.col("Close").map(lambda x: talib.KAMA(x, 30)).alias("KAMA30")
])

Answer 2

将第 15 行更改为：

return talib.KAMA(pl.col("Close"), 30)

至：

return talib.KAMA(pl_df["Close"], 30)

要完成，您还可以添加列名，来自：

pl_df2 = pl_df.select([
    pl.col("Close"),
    kama30()
])

至：

pl_df2 = pl_df.select([
    pl.col("Close"),
    pl.Series(kama30()).alias('kama30')
])

所以代码将是：

import talib
import polars as pl
import yfinance as yf

tesla = yf.Ticker('TSLA')
tesla_data = tesla.history(period="1Y")
tesla_data["Date"]=tesla_data.index
pl_df = pl.from_pandas(tesla_data[["Date", "Open", "High", "Low", "Close", "Volume"]])

# Method 1. Using ta-lib as a direct function call.
mv_kama = pl.Series(talib.KAMA(pl_df["Close"], 30)).alias('kama30')

# Method 2. Using ta-lib as Polars expression
def kama30() -> pl.Expr:
    return talib.KAMA(pl_df["Close"], 30)

pl_df2 = pl_df.select([
    pl.col("Close"),
    pl.Series(kama30()).alias('kama30')
])

print(mv_kama)
print(pl_df2)

为了保留所有数据帧列，也需要 pl.all() ：

import talib
import polars as pl
import yfinance as yf

tesla = yf.Ticker('TSLA')
tesla_data = tesla.history(period="1Y")
tesla_data["Date"]=tesla_data.index
pl_df = pl.from_pandas(tesla_data[["Date", "Open", "High", "Low", "Close", "Volume"]])

# Method 1. Using ta-lib as a direct function call.
# mv_kama = pl.Series(talib.KAMA(pl_df["Close"], 30)).alias('kama30')

# Method 2. Using ta-lib as Polars expression
def kama30() -> pl.Expr:
    return talib.KAMA(pl_df["Close"], 30)

pl_df2 = pl_df.select([
    pl.all(),
    # pl.col("Close"),
    pl.Series(kama30()).alias('kama30')
])

# print(mv_kama)
print(pl_df2)

Answer 3

要将 TA-lib 函数直接集成到 Polars 表达式中，您可以使用 Polars_talib 库。确保首先运行以下命令来安装它：

pip install polars_talib

安装后，您可以在Python脚本中使用该库：

import polars
import polars_talib as plta
import yfinance as yf

tesla = yf.Ticker('TSLA')
tesla_data = tesla.history(period="1Y")
tesla_data["Date"]=tesla_data.index
pl_df = pl.from_pandas(tesla_data[["Date", "Open", "High", "Low", "Close", "Volume"]])

pl_df2 = pl_df.select(
    pl.col("Close"),
    # Method 1: Using ta.kama from expr extension
    pl.col("Close").ta.kama(30).alias("KAMA"),
    
    # Method 2: Using kama from polars_talib directly
    # Note: Both methods are equivalent; you can choose the one you prefer
    plta.kama(pl.col("Close"), 30).alias("KAMA"),
)

在此示例中，我们使用 Polars_talib 中的 ta.kama 方法来计算“收盘”列上窗口为 30 的考夫曼自适应移动平均线。您可以根据您的具体需求修改功能和参数。

有关 Polars_talib 支持的更多详细信息和可用 TA-lib 函数列表，请参阅 github polars_ta_extension。