.\pandas-ta\tests\test_strategy.py
# 导入必要的模块
# 必须与其它测试分开运行，以确保成功运行
from multiprocessing import cpu_count # 导入获取 CPU 核心数量的函数
from time import perf_counter # 导入性能计数器
# 导入所需的配置和上下文
from .config import sample_data # 从配置文件中导入示例数据
from .context import pandas_ta # 导入上下文中的 pandas_ta 模块
from unittest import skip, skipUnless, TestCase # 导入单元测试相关的函数和类
from pandas import DataFrame # 导入 DataFrame 类
# 策略测试参数设置
cores = cpu_count() # 获取 CPU 核心数量
cumulative = False # 是否累积计算
speed_table = False # 是否生成速度表
strategy_timed = False # 是否记录策略的时间
timed = True # 是否记录时间
verbose = False # 是否冗长输出
# 测试策略方法的类
class TestStrategyMethods(TestCase):
@classmethod
def setUpClass(cls):
# 设置测试类的共享数据
cls.data = sample_data # 使用示例数据
cls.data.ta.cores = cores # 设置数据的核心数量
cls.speed_test = DataFrame() # 创建一个空的 DataFrame 用于存储速度测试结果
@classmethod
def tearDownClass(cls):
# 在测试类执行完成后的清理工作
cls.speed_test = cls.speed_test.T # 转置速度测试结果
cls.speed_test.index.name = "Test" # 设置索引名称为 "Test"
cls.speed_test.columns = ["Columns", "Seconds"] # 设置列名
if cumulative: # 如果设置了累积计算
cls.speed_test["Cum. Seconds"] = cls.speed_test["Seconds"].cumsum() # 计算累积秒数
if speed_table: # 如果设置了生成速度表
cls.speed_test.to_csv("tests/speed_test.csv") # 将速度测试结果保存为 CSV 文件
if timed: # 如果记录时间
tca = cls.speed_test['Columns'].sum() # 总列数
tcs = cls.speed_test['Seconds'].sum() # 总秒数
cps = f"[i] Total Columns / Second for All Tests: { tca / tcs:.5f} " # 计算每秒处理的列数
print("=" * len(cps)) # 打印分隔线
print(cls.speed_test) # 打印速度测试结果
print(f"[i] Cores: {cls.data.ta.cores}") # 打印核心数量
print(f"[i] Total Datapoints per run: {cls.data.shape[0]}") # 打印每次运行的数据点总数
print(f"[i] Total Columns added: {tca}") # 打印添加的总列数
print(f"[i] Total Seconds for All Tests: {tcs:.5f}") # 打印所有测试的总秒数
print(cps) # 打印每秒处理的列数
print("=" * len(cps)) # 打印分隔线
del cls.data # 删除示例数据，释放内存
def setUp(self):
# 每个测试方法执行前的设置工作
self.added_cols = 0 # 添加的列数初始化为 0
self.category = "" # 测试类别初始化为空字符串
self.init_cols = len(self.data.columns) # 记录初始列数
self.time_diff = 0 # 计算时间差初始化为 0
self.result = None # 测试结果初始化为 None
if verbose: # 如果设置了冗长输出
print() # 输出空行
if timed: # 如果记录时间
self.stime = perf_counter() # 记录开始时间
def tearDown(self):
# 每个测试方法执行后的清理工作
if timed: # 如果记录时间
self.time_diff = perf_counter() - self.stime # 计算时间差
self.added_cols = len(self.data.columns) - self.init_cols # 计算添加的列数
self.assertGreaterEqual(self.added_cols, 1) # 断言添加的列数大于等于 1
self.result = self.data[self.data.columns[-self.added_cols:]] # 获取测试结果
self.assertIsInstance(self.result, DataFrame) # 断言测试结果是 DataFrame 类型
self.data.drop(columns=self.result.columns, axis=1, inplace=True) # 在数据中删除测试结果的列
self.speed_test[self.category] = [self.added_cols, self.time_diff] # 记录测试结果到速度测试 DataFrame 中
# 测试所有策略
# @skip
def test_all(self):
self.category = "All" # 设置测试类别为 "All"
self.data.ta.strategy(verbose=verbose, timed=strategy_timed) # 调用策略函数进行测试
# 测试所有有序策略
def test_all_ordered(self):
self.category = "All" # 设置测试类别为 "All"
self.data.ta.strategy(ordered=True, verbose=verbose, timed=strategy_timed) # 调用有序策略函数进行测试
self.category = "All Ordered" # 重命名测试类别为 "All Ordered"，用于速度表显示
# 只在冗长输出模式下运行的测试方法
@skipUnless(verbose, "verbose mode only")
def test_all_strategy(self):
self.data.ta.strategy(pandas_ta.AllStrategy, verbose=verbose, timed=strategy_timed) # 调用特定策略进行测试
# 仅在 verbose 模式下才运行该测试，用于测试所有名称策略
@skipUnless(verbose, "verbose mode only")
def test_all_name_strategy(self):
self.category = "All"
# 调用 ta.strategy() 方法执行指定类别的策略，传入 verbose 和 timed 参数
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# 测试所有多参数策略
def test_all_multiparams_strategy(self):
self.category = "All"
# 调用 ta.strategy() 方法执行指定类别的策略，传入不同的参数长度
self.data.ta.strategy(self.category, length=10, verbose=verbose, timed=strategy_timed)
self.data.ta.strategy(self.category, length=50, verbose=verbose, timed=strategy_timed)
self.data.ta.strategy(self.category, fast=5, slow=10, verbose=verbose, timed=strategy_timed)
# 重命名类别以供速度表使用
self.category = "All Multiruns with diff Args"
# 测试 candles 类别策略
def test_candles_category(self):
self.category = "Candles"
# 调用 ta.strategy() 方法执行指定类别的策略，传入 verbose 和 timed 参数
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# 测试 common 类别策略
def test_common(self):
self.category = "Common"
# 调用 ta.strategy() 方法执行指定类别的策略，传入 verbose 和 timed 参数
self.data.ta.strategy(pandas_ta.CommonStrategy, verbose=verbose, timed=strategy_timed)
# 测试 cycles 类别策略
def test_cycles_category(self):
self.category = "Cycles"
# 调用 ta.strategy() 方法执行指定类别的策略，传入 verbose 和 timed 参数
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# 测试 custom A 类别策略
def test_custom_a(self):
self.category = "Custom A"
print()
print(self.category)
# 自定义指标列表
momo_bands_sma_ta = [
{"kind": "cdl_pattern", "name": "tristar"}, # 1
{"kind": "rsi"}, # 1
{"kind": "macd"}, # 3
{"kind": "sma", "length": 50}, # 1
{"kind": "sma", "length": 200 }, # 1
{"kind": "bbands", "length": 20}, # 3
{"kind": "log_return", "cumulative": True}, # 1
{"kind": "ema", "close": "CUMLOGRET_1", "length": 5, "suffix": "CLR"} # 1
]
# 创建自定义策略对象
custom = pandas_ta.Strategy(
"Commons with Cumulative Log Return EMA Chain", # name
momo_bands_sma_ta, # ta
"Common indicators with specific lengths and a chained indicator", # description
)
# 调用 ta.strategy() 方法执行自定义策略，传入 verbose 和 timed 参数
self.data.ta.strategy(custom, verbose=verbose, timed=strategy_timed)
# 断言数据列数为 15
self.assertEqual(len(self.data.columns), 15)
# 测试 custom B 类别策略
def test_custom_args_tuple(self):
self.category = "Custom B"
# 自定义指标参数列表
custom_args_ta = [
{"kind": "ema", "params": (5,)},
{"kind": "fisher", "params": (13, 7)}
]
# 创建自定义策略对象
custom = pandas_ta.Strategy(
"Custom Args Tuple",
custom_args_ta,
"Allow for easy filling in indicator arguments by argument placement."
)
# 调用 ta.strategy() 方法执行自定义策略，传入 verbose 和 timed 参数
self.data.ta.strategy(custom, verbose=verbose, timed=strategy_timed)
# 测试自定义列名元组
def test_custom_col_names_tuple(self):
# 设置类别为 "Custom C"
self.category = "Custom C"
# 自定义参数列表，包含一个字典，指定了指标类型和列名元组
custom_args_ta = [{"kind": "bbands", "col_names": ("LB", "MB", "UB", "BW", "BP")}]
# 创建自定义策略对象
custom = pandas_ta.Strategy(
"Custom Col Numbers Tuple", # 策略名称
custom_args_ta, # 自定义参数
"Allow for easy renaming of resultant columns", # 描述
)
# 应用策略到数据
self.data.ta.strategy(custom, verbose=verbose, timed=strategy_timed)
# @skip
# 测试自定义列数字元组
def test_custom_col_numbers_tuple(self):
# 设置类别为 "Custom D"
self.category = "Custom D"
# 自定义参数列表，包含一个字典，指定了指标类型和列数字元组
custom_args_ta = [{"kind": "macd", "col_numbers": (1,)}]
# 创建自定义策略对象
custom = pandas_ta.Strategy(
"Custom Col Numbers Tuple", # 策略名称
custom_args_ta, # 自定义参数
"Allow for easy selection of resultant columns", # 描述
)
# 应用策略到数据
self.data.ta.strategy(custom, verbose=verbose, timed=strategy_timed)
# @skip
# 测试自定义指标
def test_custom_a(self):
# 设置类别为 "Custom E"
self.category = "Custom E"
# 自定义参数列表，包含多个字典，指定了不同的指标类型和参数
amat_logret_ta = [
{"kind": "amat", "fast": 20, "slow": 50 }, # AMAT指标
{"kind": "log_return", "cumulative": True}, # 对数收益率
{"kind": "ema", "close": "CUMLOGRET_1", "length": 5} # 指数移动平均线
]
# 创建自定义策略对象
custom = pandas_ta.Strategy(
"AMAT Log Returns", # 策略名称
amat_logret_ta, # 自定义参数
"AMAT Log Returns", # 描述
)
# 应用策略到数据，设置ordered=True按照给定顺序执行指标计算
self.data.ta.strategy(custom, verbose=verbose, timed=strategy_timed, ordered=True)
# 添加信号列
self.data.ta.tsignals(trend=self.data["AMATe_LR_20_50_2"], append=True)
# 断言结果列数量为13
self.assertEqual(len(self.data.columns), 13)
# @skip
# 测试动量类别指标
def test_momentum_category(self):
# 设置类别为 "Momentum"
self.category = "Momentum"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skip
# 测试重叠类别指标
def test_overlap_category(self):
# 设置类别为 "Overlap"
self.category = "Overlap"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skip
# 测试性能类别指标
def test_performance_category(self):
# 设置类别为 "Performance"
self.category = "Performance"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skip
# 测试统计类别指标
def test_statistics_category(self):
# 设置类别为 "Statistics"
self.category = "Statistics"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skip
# 测试趋势类别指标
def test_trend_category(self):
# 设置类别为 "Trend"
self.category = "Trend"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skip
# 测试波动率类别指标
def test_volatility_category(self):
# 设置类别为 "Volatility"
self.category = "Volatility"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skip
# 测试交易量类别指标
def test_volume_category(self):
# 设置类别为 "Volume"
self.category = "Volume"
# 应用指定类别的策略到数据
self.data.ta.strategy(self.category, verbose=verbose, timed=strategy_timed)
# @skipUnless(verbose, "verbose mode only")
# 定义一个测试方法，用于测试不使用多进程的情况
def test_all_no_multiprocessing(self):
# 设置测试类别为"All with No Multiprocessing"
self.category = "All with No Multiprocessing"
# 保存当前核心数，并将核心数设置为0
cores = self.data.ta.cores
self.data.ta.cores = 0
# 运行策略，设置参数verbose和timed
self.data.ta.strategy(verbose=verbose, timed=strategy_timed)
# 恢复原来的核心数
self.data.ta.cores = cores
# @skipUnless(verbose, "verbose mode only")
# 定义一个测试方法，用于测试不使用多进程的自定义情况
def test_custom_no_multiprocessing(self):
# 设置测试类别为"Custom A with No Multiprocessing"
self.category = "Custom A with No Multiprocessing"
# 保存当前核心数，并将核心数设置为0
cores = self.data.ta.cores
self.data.ta.cores = 0
# 定义一组包含不同指标的策略
momo_bands_sma_ta = [
{"kind": "rsi"}, # 1
{"kind": "macd"}, # 3
{"kind": "sma", "length": 50}, # 1
{"kind": "sma", "length": 100, "col_names": "sma100"}, # 1
{"kind": "sma", "length": 200 }, # 1
{"kind": "bbands", "length": 20}, # 3
{"kind": "log_return", "cumulative": True}, # 1
{"kind": "ema", "close": "CUMLOGRET_1", "length": 5, "suffix": "CLR"}
]
# 创建自定义策略对象
custom = pandas_ta.Strategy(
"Commons with Cumulative Log Return EMA Chain", # name
momo_bands_sma_ta, # ta
"Common indicators with specific lengths and a chained indicator", # description
)
# 运行自定义策略，设置参数verbose和timed
self.data.ta.strategy(custom, verbose=verbose, timed=strategy_timed)
# 恢复原来的核心数
self.data.ta.cores = cores
.\pandas-ta\tests\test_utils.py
# 从config模块中导入sample_data变量
from .config import sample_data
# 从context模块中导入pandas_ta模块
from .context import pandas_ta
# 从unittest模块中导入skip和TestCase类
from unittest import skip, TestCase
# 从unittest.mock模块中导入patch函数
from unittest.mock import patch
# 导入numpy模块并重命名为np
import numpy as np
# 导入numpy.testing模块并重命名为npt
import numpy.testing as npt
# 从pandas模块中导入DataFrame和Series类
from pandas import DataFrame, Series
# 从pandas.api.types模块中导入is_datetime64_ns_dtype和is_datetime64tz_dtype函数
from pandas.api.types import is_datetime64_ns_dtype, is_datetime64tz_dtype
# 定义一个数据字典
data = {
"zero": [0, 0],
"a": [0, 1],
"b": [1, 0],
"c": [1, 1],
"crossed": [0, 1],
}
# 定义一个测试类TestUtilities，继承自TestCase类
class TestUtilities(TestCase):
# 类方法，设置测试类的数据
@classmethod
def setUpClass(cls):
cls.data = sample_data
# 类方法，清理测试类的数据
@classmethod
def tearDownClass(cls):
del cls.data
# 实例方法，每个测试用例执行前的初始化操作
def setUp(self):
self.crosseddf = DataFrame(data)
self.utils = pandas_ta.utils
# 实例方法，每个测试用例执行后的清理操作
def tearDown(self):
del self.crosseddf
del self.utils
# 测试用例，测试_add_prefix_suffix方法
def test__add_prefix_suffix(self):
# 测试添加前缀
result = self.data.ta.hl2(append=False, prefix="pre")
self.assertEqual(result.name, "pre_HL2")
# 测试添加后缀
result = self.data.ta.hl2(append=False, suffix="suf")
self.assertEqual(result.name, "HL2_suf")
# 测试同时添加前缀和后缀
result = self.data.ta.hl2(append=False, prefix="pre", suffix="suf")
self.assertEqual(result.name, "pre_HL2_suf")
# 测试添加数字前缀和后缀
result = self.data.ta.hl2(append=False, prefix=1, suffix=2)
self.assertEqual(result.name, "1_HL2_2")
# 测试添加前缀和后缀到MACD指标
result = self.data.ta.macd(append=False, prefix="pre", suffix="suf")
for col in result.columns:
self.assertTrue(col.startswith("pre_") and col.endswith("_suf"))
# 跳过测试用例
@skip
def test__above_below(self):
# 测试_above_below方法，判断a是否在zero上方
result = self.utils._above_below(self.crosseddf["a"], self.crosseddf["zero"], above=True)
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "a_A_zero")
npt.assert_array_equal(result, self.crosseddf["c"])
# 测试_above_below方法，判断a是否在zero下方
result = self.utils._above_below(self.crosseddf["a"], self.crosseddf["zero"], above=False)
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "a_B_zero")
npt.assert_array_equal(result, self.crosseddf["b"])
# 测试_above_below方法，判断c是否在zero上方
result = self.utils._above_below(self.crosseddf["c"], self.crosseddf["zero"], above=True)
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "c_A_zero")
npt.assert_array_equal(result, self.crosseddf["c"])
# 测试_above_below方法，判断c是否在zero下方
result = self.utils._above_below(self.crosseddf["c"], self.crosseddf["zero"], above=False)
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "c_B_zero")
npt.assert_array_equal(result, self.crosseddf["zero"])
# 测试 utils 模块中 above 函数的功能
def test_above(self):
# 调用 above 函数，传入两列数据，返回结果
result = self.utils.above(self.crosseddf["a"], self.crosseddf["zero"])
# 断言返回结果的类型为 Series
self.assertIsInstance(result, Series)
# 断言返回结果的名称为 "a_A_zero"
self.assertEqual(result.name, "a_A_zero")
# 使用 numpy.testing.assert_array_equal 方法检查返回结果是否与预期一致
npt.assert_array_equal(result, self.crosseddf["c"])
# 再次调用 above 函数，参数顺序颠倒，进行测试
result = self.utils.above(self.crosseddf["zero"], self.crosseddf["a"])
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "zero_A_a")
npt.assert_array_equal(result, self.crosseddf["b"])
# 测试 utils 模块中 above_value 函数的功能
def test_above_value(self):
# 调用 above_value 函数，传入一列数据和一个值，返回结果
result = self.utils.above_value(self.crosseddf["a"], 0)
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "a_A_0")
npt.assert_array_equal(result, self.crosseddf["c"])
# 再次调用 above_value 函数，传入不合法参数进行测试
result = self.utils.above_value(self.crosseddf["a"], self.crosseddf["zero"])
self.assertIsNone(result)
# 测试 utils 模块中 below 函数的功能
def test_below(self):
# 调用 below 函数，传入两列数据，返回结果
result = self.utils.below(self.crosseddf["zero"], self.crosseddf["a"])
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "zero_B_a")
npt.assert_array_equal(result, self.crosseddf["c"])
# 再次调用 below 函数，传入两列数据，返回结果
result = self.utils.below(self.crosseddf["zero"], self.crosseddf["a"])
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "zero_B_a")
npt.assert_array_equal(result, self.crosseddf["c"])
# 测试 utils 模块中 below_value 函数的功能
def test_below_value(self):
# 调用 below_value 函数，传入一列数据和一个值，返回结果
result = self.utils.below_value(self.crosseddf["a"], 0)
self.assertIsInstance(result, Series)
self.assertEqual(result.name, "a_B_0")
npt.assert_array_equal(result, self.crosseddf["b"])
# 再次调用 below_value 函数，传入不合法参数进行测试
result = self.utils.below_value(self.crosseddf["a"], self.crosseddf["zero"])
self.assertIsNone(result)
# 测试 utils 模块中 combination 函数的功能
def test_combination(self):
# 测试 combination 函数不传入参数时的情况
self.assertIsNotNone(self.utils.combination())
# 测试 combination 函数传入参数 n=0 时的情况
self.assertEqual(self.utils.combination(), 1)
self.assertEqual(self.utils.combination(r=-1), 1)
# 测试 combination 函数传入参数 n=10, r=4, repetition=False 时的情况
self.assertEqual(self.utils.combination(n=10, r=4, repetition=False), 210)
# 测试 combination 函数传入参数 n=10, r=4, repetition=True 时的情况
self.assertEqual(self.utils.combination(n=10, r=4, repetition=True), 715)
# 测试 utils 模块中 cross 函数的功能（above 参数为 True）
def test_cross_above(self):
# 调用 cross 函数，传入两列数据，返回结果
result = self.utils.cross(self.crosseddf["a"], self.crosseddf["b"])
self.assertIsInstance(result, Series)
# 使用 numpy.testing.assert_array_equal 方法检查返回结果是否与预期一致
npt.assert_array_equal(result, self.crosseddf["crossed"])
# 再次调用 cross 函数，传入两列数据，返回结果
result = self.utils.cross(self.crosseddf["a"], self.crosseddf["b"], above=True)
self.assertIsInstance(result, Series)
npt.assert_array_equal(result, self.crosseddf["crossed"])
# 测试 utils 模块中 cross 函数的功能（above 参数为 False）
def test_cross_below(self):
# 调用 cross 函数，传入两列数据，返回结果
result = self.utils.cross(self.crosseddf["b"], self.crosseddf["a"], above=False)
self.assertIsInstance(result, Series)
# 使用 numpy.testing.assert_array_equal 方法检查返回结果是否与预期一致
npt.assert_array_equal(result, self.crosseddf["crossed"])
# 测试 DataFrame 中日期相关函数的行为
# 调用 df_dates 函数测试，传入数据 self.data
result = self.utils.df_dates(self.data)
# 断言结果是否为 None
self.assertEqual(None, result)
# 再次调用 df_dates 函数测试，传入数据 self.data 和指定日期 "1999-11-01"
result = self.utils.df_dates(self.data, "1999-11-01")
# 断言结果行数是否为 1
self.assertEqual(1, result.shape[0])
# 再次调用 df_dates 函数测试，传入数据 self.data 和多个日期
result = self.utils.df_dates(self.data, ["1999-11-01", "2020-08-15", "2020-08-24", "2020-08-25", "2020-08-26", "2020-08-27"])
# 断言结果行数是否为 5
self.assertEqual(5, result.shape[0])
# 跳过以下测试函数
@skip
def test_df_month_to_date(self):
result = self.utils.df_month_to_date(self.data)
@skip
def test_df_quarter_to_date(self):
result = self.utils.df_quarter_to_date(self.data)
@skip
def test_df_year_to_date(self):
result = self.utils.df_year_to_date(self.data)
# 测试 Fibonacci 函数的行为
def test_fibonacci(self):
# 断言返回值类型是否为 numpy 数组
self.assertIs(type(self.utils.fibonacci(zero=True, weighted=False)), np.ndarray)
# 断言 Fibonacci 函数返回结果是否正确
npt.assert_array_equal(self.utils.fibonacci(zero=True), np.array([0, 1, 1]))
npt.assert_array_equal(self.utils.fibonacci(zero=False), np.array([1, 1]))
# 断言 Fibonacci 函数返回结果是否正确，带参数 n=0
npt.assert_array_equal(self.utils.fibonacci(n=0, zero=True, weighted=False), np.array([0]))
npt.assert_array_equal(self.utils.fibonacci(n=0, zero=False, weighted=False), np.array([1]))
# 断言 Fibonacci 函数返回结果是否正确，带参数 n=5
npt.assert_array_equal(self.utils.fibonacci(n=5, zero=True, weighted=False), np.array([0, 1, 1, 2, 3, 5]))
npt.assert_array_equal(self.utils.fibonacci(n=5, zero=False, weighted=False), np.array([1, 1, 2, 3, 5]))
# 测试带权重的 Fibonacci 函数的行为
def test_fibonacci_weighted(self):
# 断言返回值类型是否为 numpy 数组
self.assertIs(type(self.utils.fibonacci(zero=True, weighted=True)), np.ndarray)
# 断言 Fibonacci 函数返回结果是否正确，带参数 n=0
npt.assert_array_equal(self.utils.fibonacci(n=0, zero=True, weighted=True), np.array([0]))
npt.assert_array_equal(self.utils.fibonacci(n=0, zero=False, weighted=True), np.array([1]))
# 断言 Fibonacci 函数返回结果是否正确，带参数 n=5
npt.assert_allclose(self.utils.fibonacci(n=5, zero=True, weighted=True), np.array([0, 1 / 12, 1 / 12, 1 / 6, 1 / 4, 5 / 12]))
npt.assert_allclose(self.utils.fibonacci(n=5, zero=False, weighted=True), np.array([1 / 12, 1 / 12, 1 / 6, 1 / 4, 5 / 12]))
# 测试几何平均数函数的行为
def test_geometric_mean(self):
# 计算收益率并传入几何平均数函数进行测试
returns = pandas_ta.percent_return(self.data.close)
result = self.utils.geometric_mean(returns)
# 断言结果类型是否为浮点数
self.assertIsInstance(result, float)
# 传入一系列数值进行测试
result = self.utils.geometric_mean(Series([12, 14, 11, 8]))
# 断言结果类型是否为浮点数
self.assertIsInstance(result, float)
# 传入一系列数值进行测试
result = self.utils.geometric_mean(Series([100, 50, 0, 25, 0, 60]))
# 断言结果类型是否为浮点数
self.assertIsInstance(result, float)
# 传入一系列数值进行测试
series = Series([0, 1, 2, 3])
result = self.utils.geometric_mean(series)
# 断言结果类型是否为浮点数
self.assertIsInstance(result, float)
# 传入一系列数值进行测试，包括负数
result = self.utils.geometric_mean(-series)
# 断言结果类型是否为整数
self.assertIsInstance(result, int)
# 断言结果是否接近 0
self.assertAlmostEqual(result, 0)
# 测试获取时间函数
def test_get_time(self):
# 测试获取当前时间并转换为字符串
result = self.utils.get_time(to_string=True)
# 断言结果为字符串类型
self.assertIsInstance(result, str)
# 测试获取指定市场时间并转换为字符串
result = self.utils.get_time("NZSX", to_string=True)
# 断言结果字符串包含指定市场代码
self.assertTrue("NZSX" in result)
# 断言结果为字符串类型
self.assertIsInstance(result, str)
# 测试获取指定市场时间并转换为字符串
result = self.utils.get_time("SSE", to_string=True)
# 断言结果为字符串类型
self.assertIsInstance(result, str)
# 断言结果字符串包含指定市场代码
self.assertTrue("SSE" in result)
# 测试线性回归函数
def test_linear_regression(self):
# 创建示例数据
x = Series([1, 2, 3, 4, 5])
y = Series([1.8, 2.1, 2.7, 3.2, 4])
# 进行线性回归
result = self.utils.linear_regression(x, y)
# 断言结果为字典类型
self.assertIsInstance(result, dict)
# 断言字典中'a'键对应的值为浮点型
self.assertIsInstance(result["a"], float)
# 断言字典中'b'键对应的值为浮点型
self.assertIsInstance(result["b"], float)
# 断言字典中'r'键对应的值为浮点型
self.assertIsInstance(result["r"], float)
# 断言字典中't'键对应的值为浮点型
self.assertIsInstance(result["t"], float)
# 断言字典中'line'键对应的值为Series类型
self.assertIsInstance(result["line"], Series)
# 测试对数几何平均函数
def test_log_geometric_mean(self):
# 计算收益率
returns = pandas_ta.percent_return(self.data.close)
# 计算对数几何平均
result = self.utils.log_geometric_mean(returns)
# 断言结果为浮点型
self.assertIsInstance(result, float)
# 测试对数几何平均函数的其他参数
result = self.utils.log_geometric_mean(Series([12, 14, 11, 8]))
self.assertIsInstance(result, float)
result = self.utils.log_geometric_mean(Series([100, 50, 0, 25, 0, 60]))
self.assertIsInstance(result, float)
series = Series([0, 1, 2, 3])
result = self.utils.log_geometric_mean(series)
self.assertIsInstance(result, float)
result = self.utils.log_geometric_mean(-series)
# 断言结果为整型
self.assertIsInstance(result, int)
# 断言结果接近0
self.assertAlmostEqual(result, 0)
# 测试帕斯卡三角形函数
def test_pascals_triangle(self):
# 测试帕斯卡三角形的反向情况
self.assertIsNone(self.utils.pascals_triangle(inverse=True), None)
array_1 = np.array([1])
# 测试默认参数下的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(), array_1)
# 测试带权重的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(weighted=True), array_1)
# 测试带权重且反向的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(weighted=True, inverse=True), np.array([0]))
array_5 = self.utils.pascals_triangle(n=5) # or np.array([1, 5, 10, 10, 5, 1])
array_5w = array_5 / np.sum(array_5)
array_5iw = 1 - array_5w
# 测试负数行数的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(n=-5), array_5)
# 测试负数行数的带权重的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(n=-5, weighted=True), array_5w)
# 测试负数行数的带权重且反向的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(n=-5, weighted=True, inverse=True), array_5iw)
# 测试正数行数的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(n=5), array_5)
# 测试正数行数的带权重的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(n=5, weighted=True), array_5w)
# 测试正数行数的带权重且反向的帕斯卡三角形
npt.assert_array_equal(self.utils.pascals_triangle(n=5, weighted=True, inverse=True), array_5iw)
# 测试函数，验证对称三角形函数的输出是否符合预期
def test_symmetric_triangle(self):
# 验证未加权的对称三角形函数输出是否正确
npt.assert_array_equal(self.utils.symmetric_triangle(), np.array([1,1]))
# 验证加权的对称三角形函数输出是否正确
npt.assert_array_equal(self.utils.symmetric_triangle(weighted=True), np.array([0.5, 0.5])
# 验证 n=4 时对称三角形函数输出是否正确
array_4 = self.utils.symmetric_triangle(n=4) # or np.array([1, 2, 2, 1])
array_4w = array_4 / np.sum(array_4)
npt.assert_array_equal(self.utils.symmetric_triangle(n=4), array_4)
npt.assert_array_equal(self.utils.symmetric_triangle(n=4, weighted=True), array_4w)
# 验证 n=5 时对称三角形函数输出是否正确
array_5 = self.utils.symmetric_triangle(n=5) # or np.array([1, 2, 3, 2, 1])
array_5w = array_5 / np.sum(array_5)
npt.assert_array_equal(self.utils.symmetric_triangle(n=5), array_5)
npt.assert_array_equal(self.utils.symmetric_triangle(n=5, weighted=True), array_5w)
# 测试函数，验证 tal_ma 函数的输出是否符合预期
def test_tal_ma(self):
# 验证不同参数输入时 tal_ma 函数的输出是否正确
self.assertEqual(self.utils.tal_ma("sma"), 0)
self.assertEqual(self.utils.tal_ma("Sma"), 0)
self.assertEqual(self.utils.tal_ma("ema"), 1)
self.assertEqual(self.utils.tal_ma("wma"), 2)
self.assertEqual(self.utils.tal_ma("dema"), 3)
self.assertEqual(self.utils.tal_ma("tema"), 4)
self.assertEqual(self.utils.tal_ma("trima"), 5)
self.assertEqual(self.utils.tal_ma("kama"), 6)
self.assertEqual(self.utils.tal_ma("mama"), 7)
self.assertEqual(self.utils.tal_ma("t3"), 8)
# 测试函数，验证 zero 函数的输出是否符合预期
def test_zero(self):
# 验证不同参数输入时 zero 函数的输出是否正确
self.assertEqual(self.utils.zero(-0.0000000000000001), 0)
self.assertEqual(self.utils.zero(0), 0)
self.assertEqual(self.utils.zero(0.0), 0)
self.assertEqual(self.utils.zero(0.0000000000000001), 0)
self.assertNotEqual(self.utils.zero(-0.000000000000001), 0)
self.assertNotEqual(self.utils.zero(0.000000000000001), 0)
self.assertNotEqual(self.utils.zero(1), 0)
# 测试函数，验证 get_drift 函数的输出是否符合预期
def test_get_drift(self):
# 验证不同参数输入时 get_drift 函数的输出是否正确
for s in [0, None, "", [], {}]:
self.assertIsInstance(self.utils.get_drift(s), int)
self.assertEqual(self.utils.get_drift(0), 1)
self.assertEqual(self.utils.get_drift(1.1), 1)
self.assertEqual(self.utils.get_drift(-1.1), 1)
# 测试函数，验证 get_offset 函数的输出是否符合预期
def test_get_offset(self):
# 验证不同参数输入时 get_offset 函数的输出是否正确
for s in [0, None, "", [], {}]:
self.assertIsInstance(self.utils.get_offset(s), int)
self.assertEqual(self.utils.get_offset(0), 0)
self.assertEqual(self.utils.get_offset(-1.1), 0)
self.assertEqual(self.utils.get_offset(1), 1)
# 测试函数，验证 to_utc 函数的输出是否符合预期
def test_to_utc(self):
# 验证 to_utc 函数对数据的处理是否正确
result = self.utils.to_utc(self.data.copy())
self.assertTrue(is_datetime64_ns_dtype(result.index))
self.assertTrue(is_datetime64tz_dtype(result.index))
# 测试计算给定数据的总时间
def test_total_time(self):
# 计算总时间，默认单位为年
result = self.utils.total_time(self.data)
# 断言总时间为约30.18年
self.assertEqual(30.182539682539684, result)
# 计算总时间，单位为月
result = self.utils.total_time(self.data, "months")
# 断言总时间为约250.06个月
self.assertEqual(250.05753361606995, result)
# 计算总时间，单位为周
result = self.utils.total_time(self.data, "weeks")
# 断言总时间为约1086.57周
self.assertEqual(1086.5714285714287, result)
# 计算总时间，单位为天
result = self.utils.total_time(self.data, "days")
# 断言总时间为7606天
self.assertEqual(7606, result)
# 计算总时间，单位为小时
result = self.utils.total_time(self.data, "hours")
# 断言总时间为182544小时
self.assertEqual(182544, result)
# 计算总时间，单位为分钟
result = self.utils.total_time(self.data, "minutes")
# 断言总时间为10952640分钟
self.assertEqual(10952640.0, result)
# 计算总时间，单位为秒
result = self.utils.total_time(self.data, "seconds")
# 断言总时间为657158400秒
self.assertEqual(657158400.0, result)
# 测试 pandas_ta 库的版本
def test_version(self):
# 获取 pandas_ta 库的版本号
result = pandas_ta.version
# 断言版本号为字符串类型
self.assertIsInstance(result, str)
# 打印 pandas_ta 库的版本信息
print(f"\nPandas TA v{result}")
.\pandas-ta\tests\test_utils_metrics.py
# 导入 unittest 模块中的 skip 和 TestCase 类
from unittest import skip, TestCase
# 导入 pandas 模块中的 DataFrame 类
from pandas import DataFrame
# 从当前目录下的 config 模块中导入 sample_data 变量
from .config import sample_data
# 从当前目录下的 context 模块中导入 pandas_ta 模块
from .context import pandas_ta
# 定义测试类 TestUtilityMetrics，继承自 TestCase 类
class TestUtilityMetrics(TestCase):
# 类方法，设置测试所需的数据
@classmethod
def setUpClass(cls):
cls.data = sample_data
cls.close = cls.data["close"]
# 计算价格序列的百分比收益率，不累积
cls.pctret = pandas_ta.percent_return(cls.close, cumulative=False)
# 计算价格序列的对数收益率，不累积
cls.logret = pandas_ta.percent_return(cls.close, cumulative=False)
# 类方法，清理测试所需的数据
@classmethod
def tearDownClass(cls):
del cls.data
del cls.pctret
del cls.logret
# 设置测试用例的前置操作
def setUp(self): pass
# 设置测试用例的后置操作
def tearDown(self): pass
# 测试计算 CAGR 的函数
def test_cagr(self):
# 调用 utils 模块中的 cagr 函数计算 CAGR
result = pandas_ta.utils.cagr(self.data.close)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于 0
self.assertGreater(result, 0)
# 测试计算 Calmar 比率的函数
def test_calmar_ratio(self):
# 调用 calmar_ratio 函数计算 Calmar 比率
result = pandas_ta.calmar_ratio(self.close)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 传入参数 years=0 调用 calmar_ratio 函数
result = pandas_ta.calmar_ratio(self.close, years=0)
# 断言返回结果为 None
self.assertIsNone(result)
# 传入参数 years=-2 调用 calmar_ratio 函数
result = pandas_ta.calmar_ratio(self.close, years=-2)
# 断言返回结果为 None
self.assertIsNone(result)
# 测试计算下行偏差的函数
def test_downside_deviation(self):
# 调用 downside_deviation 函数计算下行偏差
result = pandas_ta.downside_deviation(self.pctret)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 调用 downside_deviation 函数计算下行偏差
result = pandas_ta.downside_deviation(self.logret)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 测试计算回撤的函数
def test_drawdown(self):
# 调用 drawdown 函数计算回撤
result = pandas_ta.drawdown(self.pctret)
# 断言返回结果类型为 DataFrame
self.assertIsInstance(result, DataFrame)
# 断言结果的名称为 "DD"
self.assertEqual(result.name, "DD")
# 调用 drawdown 函数计算回撤
result = pandas_ta.drawdown(self.logret)
# 断言返回结果类型为 DataFrame
self.assertIsInstance(result, DataFrame)
# 断言结果的名称为 "DD"
self.assertEqual(result.name, "DD")
# 测试计算 Jensen's Alpha 的函数
def test_jensens_alpha(self):
# 从百分比收益率中随机抽取与收盘价序列长度相同的样本作为基准收益率
bench_return = self.pctret.sample(n=self.close.shape[0], random_state=1)
# 调用 jensens_alpha 函数计算 Jensen's Alpha
result = pandas_ta.jensens_alpha(self.close, bench_return)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 测试计算对数最大回撤的函数
def test_log_max_drawdown(self):
# 调用 log_max_drawdown 函数计算对数最大回撤
result = pandas_ta.log_max_drawdown(self.close)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 测试计算最大回撤的函数
def test_max_drawdown(self):
# 调用 max_drawdown 函数计算最大回撤
result = pandas_ta.max_drawdown(self.close)
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 传入参数 method="percent" 调用 max_drawdown 函数
result = pandas_ta.max_drawdown(self.close, method="percent")
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 传入参数 method="log" 调用 max_drawdown 函数
result = pandas_ta.max_drawdown(self.close, method="log")
# 断言返回结果类型为浮点数
self.assertIsInstance(result, float)
# 断言计算结果大于等于 0
self.assertGreaterEqual(result, 0)
# 传入参数 all=True 调用 max_drawdown 函数
result = pandas_ta.max_drawdown(self.close, all=True)
# 断言返回结果类型为字典
self.assertIsInstance(result, dict)
# 断言字
# 测试 optimal_leverage 函数，计算最佳杠杆倍数
def test_optimal_leverage(self):
# 调用 optimal_leverage 函数，计算默认参数情况下的最佳杠杆倍数
result = pandas_ta.optimal_leverage(self.close)
# 断言结果为整数类型
self.assertIsInstance(result, int)
# 调用 optimal_leverage 函数，计算开启日志情况下的最佳杠杆倍数
result = pandas_ta.optimal_leverage(self.close, log=True)
# 断言结果为整数类型
self.assertIsInstance(result, int)
# 测试 pure_profit_score 函数，计算纯利得分
def test_pure_profit_score(self):
# 调用 pure_profit_score 函数，计算纯利得分
result = pandas_ta.pure_profit_score(self.close)
# 断言结果大于等于0
self.assertGreaterEqual(result, 0)
# 测试 sharpe_ratio 函数，计算夏普比率
def test_sharpe_ratio(self):
# 调用 sharpe_ratio 函数，计算夏普比率
result = pandas_ta.sharpe_ratio(self.close)
# 断言结果为浮点数类型
self.assertIsInstance(result, float)
# 断言结果大于等于0
self.assertGreaterEqual(result, 0)
# 测试 sortino_ratio 函数，计算索提诺比率
def test_sortino_ratio(self):
# 调用 sortino_ratio 函数，计算索提诺比率
result = pandas_ta.sortino_ratio(self.close)
# 断言结果为浮点数类型
self.assertIsInstance(result, float)
# 断言结果大于等于0
self.assertGreaterEqual(result, 0)
# 测试 volatility 函数，计算波动率
def test_volatility(self):
# 计算收益率
returns_ = pandas_ta.percent_return(self.close)
# 调用 volatility 函数，计算波动率，返回结果包含收益率信息
result = pandas_ta.utils.volatility(returns_, returns=True)
# 断言结果为浮点数类型
self.assertIsInstance(result, float)
# 断言结果大于等于0
self.assertGreaterEqual(result, 0)
# 遍历不同时间段
for tf in ["years", "months", "weeks", "days", "hours", "minutes", "seconds"]:
# 调用 volatility 函数，计算指定时间段的波动率
result = pandas_ta.utils.volatility(self.close, tf)
# 使用子测试，传入时间段参数
with self.subTest(tf=tf):
# 断言结果为浮点数类型
self.assertIsInstance(result, float)
# 断言结果大于等于0
self.assertGreaterEqual(result, 0)
