Chan/chanlun/tests/test_ta_compat.py

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