"""Pin chanlun.indicators.ta to TA-Lib's output, bar for bar. These indicators feed the Chan structure builders, so a one-bar shift in the warm-up or a different smoothing seed silently changes every downstream bi/seg/zs. Equality against the reference implementation is the only check that catches that. Skipped when talib is unavailable — which is the point of the replacement, so the suite still has to pass without it. Run in an environment that has talib whenever chanlun/indicators/ta.py changes. """ from __future__ import annotations import sys import unittest from pathlib import Path import numpy as np import pandas as pd sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from chanlun.indicators import ta # noqa: E402 from chanlun.pipeline.resample import resample_to_interval # noqa: E402 try: import talib.abstract as reference except ImportError: # pragma: no cover reference = None requires_talib = unittest.skipIf(reference is None, "talib not installed") def make_ohlcv(n: int = 900, seed: int = 7) -> pd.DataFrame: """Random walk with enough range for BBANDS(365) and EMA(208) to warm up.""" rng = np.random.default_rng(seed) close = 100.0 + np.cumsum(rng.normal(0.0, 1.0, n)) spread = np.abs(rng.normal(0.0, 0.6, n)) + 0.05 high = close + spread low = close - spread open_ = np.concatenate([[close[0]], close[:-1]]) return pd.DataFrame( { "date": pd.date_range("2024-01-01", periods=n, freq="1min", tz="UTC"), "open": open_, "high": np.maximum.reduce([high, open_, close]), "low": np.minimum.reduce([low, open_, close]), "close": close, "volume": rng.uniform(1.0, 100.0, n), } ) class TAEquivalence(unittest.TestCase): def setUp(self) -> None: self.df = make_ohlcv() def assertSameSeries(self, got, expected, label: str) -> None: g = np.asarray(got, dtype=float) e = np.asarray(expected, dtype=float) self.assertEqual(g.shape, e.shape, f"{label}: shape") np.testing.assert_array_equal( np.isnan(g), np.isnan(e), err_msg=f"{label}: NaN warm-up differs" ) mask = ~np.isnan(e) np.testing.assert_allclose( g[mask], e[mask], rtol=1e-9, atol=1e-8, err_msg=f"{label}: values differ" ) @requires_talib def test_sma(self) -> None: for period in (5, 20, 90, 250): self.assertSameSeries( ta.SMA(self.df, timeperiod=period), reference.SMA(self.df, timeperiod=period), f"SMA({period})", ) @requires_talib def test_ma(self) -> None: for period in (5, 10, 250): self.assertSameSeries( ta.MA(self.df, timeperiod=period), reference.MA(self.df, timeperiod=period), f"MA({period})", ) @requires_talib def test_ema(self) -> None: for period in (5, 7, 10, 13, 24, 26, 30, 52, 104, 156, 208): self.assertSameSeries( ta.EMA(self.df, timeperiod=period), reference.EMA(self.df, timeperiod=period), f"EMA({period})", ) @requires_talib def test_rsi(self) -> None: for period in (7, 14, 21): self.assertSameSeries( ta.RSI(self.df, timeperiod=period), reference.RSI(self.df, timeperiod=period), f"RSI({period})", ) @requires_talib def test_atr(self) -> None: for period in (7, 14, 30): self.assertSameSeries( ta.ATR(self.df, timeperiod=period), reference.ATR(self.df, timeperiod=period), f"ATR({period})", ) @requires_talib def test_macd(self) -> None: for fast, slow, signal in ((12, 26, 9), (26, 52, 9), (5, 35, 5)): got = ta.MACD(self.df, fastperiod=fast, slowperiod=slow, signalperiod=signal) exp = reference.MACD(self.df, fastperiod=fast, slowperiod=slow, signalperiod=signal) for col in ("macd", "macdsignal", "macdhist"): self.assertSameSeries(got[col], exp[col], f"MACD({fast},{slow},{signal}).{col}") @requires_talib def test_bbands(self) -> None: cases = ( (365, 3.0, 3.0), (120, 3.0, 3.0), (41, 2.3, 2.3), (41, 2.0, 2.0), (26, 3.0, 3.0), (20, 2.0, 2.0), (14, 2.0, 2.0), ) for period, up, dn in cases: got = ta.BBANDS(self.df, timeperiod=period, nbdevup=up, nbdevdn=dn, matype=0) exp = reference.BBANDS(self.df, timeperiod=period, nbdevup=up, nbdevdn=dn, matype=0) for col in ("upperband", "middleband", "lowerband"): self.assertSameSeries(got[col], exp[col], f"BBANDS({period},{up},{dn}).{col}") @requires_talib def test_bbands_is_more_accurate_than_talib_on_tiny_windows(self) -> None: """A deliberate divergence, documented so nobody "fixes" it back. TA-Lib derives the variance from sumsq/n - mean**2, which cancels catastrophically when the window is short and prices are far from zero; at timeperiod=2 it drifts ~1e-6. Rolling std is accurate there, so the two disagree. No timeperiod below 14 is used in this codebase, and the periods that are used agree to ~1e-10 (covered by test_bbands). """ got = ta.BBANDS(self.df, timeperiod=2, nbdevup=1.0, nbdevdn=1.0, matype=0)["upperband"] exp = reference.BBANDS(self.df, timeperiod=2, nbdevup=1.0, nbdevdn=1.0, matype=0)["upperband"] window = self.df["close"].rolling(2) truth = window.mean() + window.std(ddof=0) ours = np.nanmax(np.abs((got - truth).to_numpy())) theirs = np.nanmax(np.abs((exp - truth).to_numpy())) self.assertLess(ours, 1e-9) self.assertLess(ours, theirs) @requires_talib def test_matches_on_real_price_scale(self) -> None: """Guard against tolerances that only hold near 100.""" df = self.df.copy() for col in ("open", "high", "low", "close"): df[col] *= 900.0 self.assertSameSeries( ta.ATR(df, timeperiod=14), reference.ATR(df, timeperiod=14), "ATR@scale" ) self.assertSameSeries( ta.RSI(df, timeperiod=14), reference.RSI(df, timeperiod=14), "RSI@scale" ) class ResampleEquivalence(unittest.TestCase): @unittest.skipIf( __import__("importlib").util.find_spec("technical") is None, "technical not installed", ) def test_matches_technical(self) -> None: from technical.util import resample_to_interval as ref_resample df = make_ohlcv(600) for interval in (5, 15, 60, "5m", "1h"): got = resample_to_interval(df, interval) exp = ref_resample(df, interval) pd.testing.assert_frame_equal(got, exp, check_exact=False, rtol=1e-12) def test_shapes_without_reference(self) -> None: df = make_ohlcv(120) out = resample_to_interval(df, 5) self.assertEqual(list(out.columns), ["date", "open", "high", "low", "close", "volume"]) self.assertLessEqual(len(out), 120 // 5 + 1) self.assertTrue((out["high"] >= out["low"]).all()) if __name__ == "__main__": unittest.main()