""" 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"])