395 lines
12 KiB
Python
395 lines
12 KiB
Python
"""
|
|
Equivalence tests for fast polars operators vs vnpy original implementations.
|
|
|
|
TDD approach: Tests define exact behavior, fast implementations must match.
|
|
Each test covers: NaN handling, ties, edge cases, multi-symbol correctness.
|
|
"""
|
|
|
|
import sys
|
|
import numpy as np
|
|
import polars as pl
|
|
import pytest
|
|
|
|
# Inject vnpy source path
|
|
_VNPY_SRC = "/Users/chufeng/.openclaw/sanguo_projects/sanguo_vnpy_v2/vnpy_v4.4.0"
|
|
if _VNPY_SRC not in sys.path:
|
|
sys.path.insert(0, _VNPY_SRC)
|
|
|
|
from vnpy.alpha.dataset.utility import DataProxy
|
|
from vnpy.alpha.dataset import ts_function
|
|
|
|
from sanguo_factor.fast_ops import (
|
|
register_fast_ops,
|
|
fast_ts_rank,
|
|
fast_ts_corr,
|
|
fast_ts_cov,
|
|
fast_ts_decay_linear,
|
|
fast_ts_slope,
|
|
fast_ts_rsquare,
|
|
fast_ts_resi,
|
|
fast_ts_quantile,
|
|
)
|
|
|
|
|
|
def create_dataproxy(values: np.ndarray, symbols: list[str] = ["000001.SZ"]) -> DataProxy:
|
|
"""Helper: Create DataProxy from numpy array."""
|
|
n = len(values)
|
|
# Create datetime sequence with proper dtype
|
|
import pandas as pd
|
|
datetimes_pd = pd.date_range("2024-01-01", periods=n, freq="1d")
|
|
datetimes = pl.Series(datetimes_pd)
|
|
|
|
# Repeat for each symbol
|
|
datetimes_list = []
|
|
symbols_list = []
|
|
values_list = []
|
|
|
|
for symbol in symbols:
|
|
datetimes_list.extend(datetimes)
|
|
symbols_list.extend([symbol] * n)
|
|
values_list.extend(values)
|
|
|
|
df = pl.DataFrame({
|
|
"datetime": datetimes_list,
|
|
"vt_symbol": symbols_list,
|
|
"data": values_list,
|
|
})
|
|
|
|
return DataProxy(df)
|
|
|
|
|
|
def assert_dataproxy_equal(
|
|
result1: DataProxy,
|
|
result2: DataProxy,
|
|
rtol: float = 1e-9,
|
|
atol: float = 1e-9,
|
|
check_nan: bool = True,
|
|
):
|
|
"""Assert two DataProxy objects are equal."""
|
|
df1 = result1.df.sort("datetime", "vt_symbol")
|
|
df2 = result2.df.sort("datetime", "vt_symbol")
|
|
|
|
# Check column structure
|
|
assert list(df1.columns) == list(df2.columns)
|
|
|
|
# Check data column values
|
|
data1 = df1["data"].to_numpy()
|
|
data2 = df2["data"].to_numpy()
|
|
|
|
if check_nan:
|
|
np.testing.assert_allclose(data1, data2, rtol=rtol, atol=atol)
|
|
else:
|
|
mask = ~(np.isnan(data1) | np.isnan(data2))
|
|
np.testing.assert_allclose(data1[mask], data2[mask], rtol=rtol, atol=atol)
|
|
|
|
|
|
class TestFastTsRank:
|
|
"""Test fast_ts_rank equivalence."""
|
|
|
|
def test_random_unique_values(self):
|
|
"""Exact match on random unique values."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 100 + 500 # Large range, unlikely ties
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_rank(feature, window=10)
|
|
fast = fast_ts_rank(feature, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-9)
|
|
|
|
def test_with_ties_approximate(self):
|
|
"""With ties, rank may vary but must be in [0, 1]."""
|
|
values = np.array([10.0, 10.0, 10.0, 5.0, 15.0] * 20)
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_rank(feature, window=5)
|
|
fast = fast_ts_rank(feature, window=5)
|
|
|
|
# Check range
|
|
df = fast.df
|
|
assert (df["data"] >= 0).all()
|
|
assert (df["data"] <= 1).all()
|
|
|
|
def test_with_nan(self):
|
|
"""Handle NaN correctly."""
|
|
values = np.array([1.0, np.nan, 3.0, np.nan, 5.0] * 10)
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_rank(feature, window=5)
|
|
fast = fast_ts_rank(feature, window=5)
|
|
|
|
assert_dataproxy_equal(original, fast, check_nan=False)
|
|
|
|
|
|
class TestFastTsCorr:
|
|
"""Test fast_ts_corr equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values1 = np.random.randn(100) * 10 + 50
|
|
values2 = np.random.randn(100) * 10 + 50
|
|
|
|
feature1 = create_dataproxy(values1)
|
|
feature2 = create_dataproxy(values2)
|
|
|
|
original = ts_function.ts_corr(feature1, feature2, window=10)
|
|
fast = fast_ts_corr(feature1, feature2, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-7) # Looser tolerance for correlation
|
|
|
|
def test_perfect_correlation(self):
|
|
"""Perfect correlation should be exactly 1.0."""
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature1 = create_dataproxy(values)
|
|
feature2 = create_dataproxy(values * 2 + 10) # Linear transformation
|
|
|
|
original = ts_function.ts_corr(feature1, feature2, window=10)
|
|
fast = fast_ts_corr(feature1, feature2, window=10)
|
|
|
|
# Should be exactly 1.0 where correlation is defined
|
|
df_fast = fast.df
|
|
assert (df_fast.filter(pl.col("data").is_not_null())["data"] - 1.0).abs().max() < 1e-6
|
|
|
|
|
|
class TestFastTsCov:
|
|
"""Test fast_ts_cov equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values1 = np.random.randn(100) * 10 + 50
|
|
values2 = np.random.randn(100) * 10 + 50
|
|
|
|
feature1 = create_dataproxy(values1)
|
|
feature2 = create_dataproxy(values2)
|
|
|
|
original = ts_function.ts_cov(feature1, feature2, window=10)
|
|
fast = fast_ts_cov(feature1, feature2, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-7)
|
|
|
|
|
|
class TestFastTsDecayLinear:
|
|
"""Test fast_ts_decay_linear equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_decay_linear(feature, window=10)
|
|
fast = fast_ts_decay_linear(feature, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-9)
|
|
|
|
|
|
class TestFastTsSlope:
|
|
"""Test fast_ts_slope equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_slope(feature, window=10)
|
|
fast = fast_ts_slope(feature, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-7)
|
|
|
|
def test_linear_trend(self):
|
|
"""Perfect linear trend should have exact slope."""
|
|
x = np.arange(100)
|
|
values = 2.5 * x + 10.0
|
|
feature = create_dataproxy(values)
|
|
|
|
result = fast_ts_slope(feature, window=20)
|
|
|
|
# Slope should be approximately 2.5
|
|
df = result.df
|
|
slopes = df.filter(pl.col("data").is_not_null())["data"].to_numpy()
|
|
assert np.abs(slopes[50] - 2.5) < 0.1 # Middle of window
|
|
|
|
|
|
class TestFastTsRsquare:
|
|
"""Test fast_ts_rsquare equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_rsquare(feature, window=10)
|
|
fast = fast_ts_rsquare(feature, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-7)
|
|
|
|
def test_linear_trend(self):
|
|
"""Perfect linear trend should have R² = 1."""
|
|
x = np.arange(100)
|
|
values = 2.5 * x + 10.0
|
|
feature = create_dataproxy(values)
|
|
|
|
result = fast_ts_rsquare(feature, window=20)
|
|
|
|
df = result.df
|
|
r2_values = df.filter(pl.col("data").is_not_null())["data"].to_numpy()
|
|
assert r2_values[50] > 0.99 # Should be near 1.0
|
|
|
|
|
|
class TestFastTsResi:
|
|
"""Test fast_ts_resi equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_resi(feature, window=10)
|
|
fast = fast_ts_resi(feature, window=10)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-7)
|
|
|
|
|
|
class TestFastTsQuantile:
|
|
"""Test fast_ts_quantile equivalence."""
|
|
|
|
def test_random_data(self):
|
|
"""Exact match on random data."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
original = ts_function.ts_quantile(feature, window=10, quantile=0.5)
|
|
fast = fast_ts_quantile(feature, window=10, quantile=0.5)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-9)
|
|
|
|
def test_multiple_quantiles(self):
|
|
"""Test various quantile values."""
|
|
np.random.seed(42)
|
|
values = np.random.randn(100) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
for q in [0.25, 0.5, 0.75, 0.9]:
|
|
original = ts_function.ts_quantile(feature, window=10, quantile=q)
|
|
fast = fast_ts_quantile(feature, window=10, quantile=q)
|
|
|
|
assert_dataproxy_equal(original, fast, rtol=1e-9)
|
|
|
|
|
|
class TestRegistrationAndIntegration:
|
|
"""Test registration and vnpy integration."""
|
|
|
|
def test_registration_idempotent(self):
|
|
"""Registration should be idempotent."""
|
|
overrides1 = register_fast_ops()
|
|
overrides2 = register_fast_ops()
|
|
|
|
assert overrides1 == overrides2
|
|
assert len(overrides1) == 8
|
|
|
|
def test_expression_override(self):
|
|
"""Test that registered functions override vnpy defaults."""
|
|
from vnpy.alpha.dataset.utility import EXPRESSION_FUNCTIONS, calculate_by_expression
|
|
|
|
# Register fast ops
|
|
register_fast_ops()
|
|
|
|
# Check that fast ops are registered
|
|
assert "ts_rank" in EXPRESSION_FUNCTIONS
|
|
assert "ts_corr" in EXPRESSION_FUNCTIONS
|
|
|
|
# Test through calculate_by_expression
|
|
np.random.seed(42)
|
|
values = np.random.randn(20) * 10 + 50
|
|
feature = create_dataproxy(values)
|
|
|
|
# Test ts_rank through expression (pass DataFrame, not DataProxy)
|
|
result = calculate_by_expression(feature.df, "ts_rank(data, 5)")
|
|
|
|
# Should not crash and should return valid DataFrame result
|
|
assert result is not None
|
|
# calculate_by_expression returns DataFrame, not DataProxy
|
|
assert hasattr(result, 'columns') # It's a DataFrame
|
|
df = result # result is already a DataFrame
|
|
assert (df["data"] >= 0).all()
|
|
assert (df["data"] <= 1).all()
|
|
|
|
def test_fast_ops_registered_count(self):
|
|
"""Verify all expected operators are registered."""
|
|
overrides = register_fast_ops()
|
|
|
|
expected = {
|
|
"ts_rank", "ts_corr", "ts_cov",
|
|
"ts_decay_linear", "ts_slope", "ts_rsquare", "ts_resi", "ts_quantile"
|
|
}
|
|
|
|
assert set(overrides) == expected
|
|
|
|
|
|
class TestEdgeCases:
|
|
"""Test edge cases and boundary conditions."""
|
|
|
|
def test_short_series(self):
|
|
"""Series shorter than window should not crash."""
|
|
values = np.array([1.0, 2.0, 3.0])
|
|
feature = create_dataproxy(values)
|
|
|
|
# Should not crash
|
|
for fast_func in [fast_ts_rank, fast_ts_decay_linear, fast_ts_slope]:
|
|
result = fast_func(feature, window=10)
|
|
assert result is not None
|
|
|
|
def test_all_nan(self):
|
|
"""All NaN values should be handled."""
|
|
values = np.array([np.nan] * 20)
|
|
feature = create_dataproxy(values)
|
|
|
|
for fast_func in [fast_ts_rank, fast_ts_decay_linear, fast_ts_slope]:
|
|
result = fast_func(feature, window=5)
|
|
assert result is not None
|
|
|
|
def test_constant_values(self):
|
|
"""Constant values should be handled."""
|
|
values = np.array([5.0] * 20)
|
|
feature = create_dataproxy(values)
|
|
|
|
for fast_func in [fast_ts_rank, fast_ts_decay_linear, fast_ts_slope]:
|
|
result = fast_func(feature, window=5)
|
|
assert result is not None
|
|
|
|
def test_single_value(self):
|
|
"""Single value should not crash."""
|
|
values = np.array([5.0])
|
|
feature = create_dataproxy(values)
|
|
|
|
for fast_func in [fast_ts_rank, fast_ts_decay_linear, fast_ts_slope]:
|
|
result = fast_func(feature, window=5)
|
|
assert result is not None
|
|
|
|
def test_shift_unroll_speed_guard(self):
|
|
"""Performance guard to prevent regression to rolling_map."""
|
|
import time
|
|
n = 200_000
|
|
df = pl.DataFrame({
|
|
"vt_symbol": ["A"] * n,
|
|
"datetime": list(range(n)),
|
|
"data": [((i * 37) % 997) / 997 for i in range(n)]
|
|
})
|
|
feature = DataProxy(df)
|
|
|
|
t0 = time.time()
|
|
fast_ts_rank(feature, 20)
|
|
elapsed = time.time() - t0
|
|
|
|
# Should complete in under 10 seconds for 200k rows
|
|
assert elapsed < 10, f"Performance regression: {elapsed:.2f}s > 10s"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"]) |