"""TACHandler: a qlib DataHandlerLP that builds datasets from the TradeAC lake. This is the "custom DataHandler" entry point (Option B): the handler is referenced from the workflow yaml's ``dataset.handler`` and reads OHLCV + pre-computed ta-lib features straight from the lake parquet files through ``QLibDataLoader`` + the tac_qlib feature provider. The standard qlib processor pipeline (``infer_processors`` / ``learn_processors``) still runs on top, so existing recipes such as ``DropnaLabel``, ``CSZScoreNorm`` or ``RobustZScoreNorm`` keep working unchanged. """ from __future__ import annotations import os from inspect import getfullargspec from typing import List, Optional, Tuple, Union import numpy as np import pandas as pd from qlib.data.dataset import processor as processor_module from qlib.data.dataset.handler import DataHandlerLP from qlib.utils import get_callable_kwargs from ...data.config import ( LakeConfig, timeframe_for_freq, FEATURE_FAMILIES, NON_FEATURE_COLUMNS, ) DEFAULT_INFER_PROCESSORS = [ {"class": "DropAllNaN", "kwargs": {}}, {"class": "ProcessInf", "kwargs": {}}, {"class": "ZScoreNorm", "kwargs": {}}, {"class": "Fillna", "kwargs": {}}, ] DEFAULT_LEARN_PROCESSORS = [ {"class": "DropnaLabel"}, {"class": "CSZScoreNorm", "kwargs": {"fields_group": "label"}}, ] #: always include raw OHLCV; ta-lib columns are discovered from the lake and appended. RAW_FEATURE_FIELDS = ("$open", "$high", "$low", "$close", "$vwap", "$volume") DEFAULT_LABEL = "Ref($close,-2)/Ref($close,-1)-1" def check_transform_proc(proc_l, fit_start_time, fit_end_time): """Port of ``qlib.contrib.data.handler.check_transform_proc`` (inject fit window into procs).""" new_l = [] for p in proc_l: if not isinstance(p, processor_module.Processor): klass, pkwargs = get_callable_kwargs(p, processor_module) args = getfullargspec(klass).args if "fit_start_time" in args and "fit_end_time" in args: assert fit_start_time is not None and fit_end_time is not None, ( "Make sure `fit_start_time` and `fit_end_time` are not None." ) pkwargs.update({"fit_start_time": fit_start_time, "fit_end_time": fit_end_time}) proc_config = {"class": klass.__name__, "kwargs": pkwargs} if isinstance(p, dict) and "module_path" in p: proc_config["module_path"] = p["module_path"] new_l.append(proc_config) else: new_l.append(p) return new_l def get_common_feature_fields(lake_root=None, market="US", timeframe="1d") -> List[str]: """Discover feature columns present in *every* feature file of the lake. Walks the `family=ta|sp` partition layout (plus any legacy flat files). TA and SP columns are disjoint by construction, so the common set is computed per family (columns shared by all symbol files of that family), then the per-family results are unioned. Returns sorted field names (without the ``$`` prefix). Empty if no features are persisted. """ cfg = LakeConfig(lake_root, market) feat_dir = cfg.features_dir(timeframe) if not feat_dir.exists(): return [] import pyarrow.parquet as pq def _family_common(fam_dir: Path) -> set: common = None for p in sorted(fam_dir.glob("symbol=*.parquet")): try: cols = set(pq.read_schema(p).names) - set(NON_FEATURE_COLUMNS) except Exception: # pragma: no cover - skip unreadable files continue common = cols if common is None else (common & cols) if not common: break return common or set() common: set = set() # family tier: features/market=*/timeframe=*/family=*/symbol=*.parquet for fam in FEATURE_FAMILIES: fam_dir = feat_dir / f"family={fam}" if fam_dir.is_dir(): common |= _family_common(fam_dir) # legacy flat: features/market=*/timeframe=*/symbol=*.parquet if (feat_dir / "family=ta").exists() or (feat_dir / "family=sp").exists(): pass # family layout already covered else: common |= _family_common(feat_dir) return sorted(common) class DropAllNaN(processor_module.Processor): """Drop feature columns that are all-NaN over the fit window. The lake can hold fully-empty indicator columns (e.g. a ta-lib output that was NaN from the start). Such columns carry no learnable signal and make ``ZScoreNorm.fit`` warn on empty slices, so we drop them before any other processor runs. The drop set is fixed on the fit window once (during ``fit``), then applied consistently to every segment so train/valid/test keep identical feature columns. """ def __init__(self, fit_start_time=None, fit_end_time=None): self.fit_start_time = fit_start_time self.fit_end_time = fit_end_time self.cols_to_drop = [] def fit(self, df=None): if df is None or len(df) == 0: return self window = df if self.fit_start_time is not None and self.fit_end_time is not None: try: from qlib.data.dataset.utils import fetch_df_by_index window = fetch_df_by_index( df, slice(self.fit_start_time, self.fit_end_time), level="datetime" ) except Exception: # pragma: no cover - defensive window = df if len(window) == 0: return self self.cols_to_drop = [c for c in window.columns if window[c].isna().all()] return self def __call__(self, df): if self.cols_to_drop: return df.drop(columns=self.cols_to_drop, errors="ignore") return df class BenchResidual(processor_module.Processor): """Subtract a benchmark instrument's forward return from the label, per datetime. Turns the training target from an absolute-return rank into a *residual* rank: ``r_i - r_bench`` is ranked cross-sectionally by the downstream ``CSRankNorm`` / ``CSZScoreNorm`` processors instead of ``r_i`` alone. Must be inserted BEFORE any per-date normalization so the ranking itself is computed on residual returns (ordering flips exactly where the benchmark trends). Stateless: ``fit`` is a no-op and the benchmark forward return is recomputed from the lake parquet on first ``__call__``. Rows whose benchmark value is missing are left untouched. Accepts ``fit_start_time``/``fit_end_time`` (ignored) so ``check_transform_proc`` can inject the fit window uniformly. NOTE: under any cross-sectional normalization downstream (``CSRankNorm`` / ``CSZScoreNorm``) this processor is a mathematical no-op: subtracting the same per-date constant preserves ranks, and z-scoring absorbs constant shifts. Use ``BenchBetaResidual`` for a target that actually reorders. """ def __init__( self, benchmark="SPY", fields_group="label", lake_root=None, market="US", timeframe=None, freq="day", fit_start_time=None, fit_end_time=None, ): self.benchmark = benchmark self.fields_group = fields_group self.lake_root = lake_root self.market = market self.timeframe = timeframe or timeframe_for_freq(freq) self.fit_start_time = fit_start_time self.fit_end_time = fit_end_time self._bench_label = None def _load_bench_label(self): if self._bench_label is not None: return self._bench_label cfg = LakeConfig(self.lake_root, self.market) p = cfg.bar_path(self.timeframe, self.benchmark) if not p.exists(): raise FileNotFoundError(f"BenchResidual: benchmark bar file not found: {p}") df = pd.read_parquet(p) s = pd.Series(df["c"].astype(float).values, index=pd.to_datetime(df["t"])).sort_index() s.index = s.index.normalize() # mirror Ref($close,-6)/Ref($close,-1)-1 on the benchmark's own calendar bench_label = s.shift(-6) / s.shift(-1) - 1 self._bench_label = bench_label[~bench_label.index.duplicated(keep="last")] return self._bench_label def fit(self, df=None): return self def __call__(self, df): bl = self._load_bench_label() cols = processor_module.get_group_columns(df, self.fields_group) dt = df.index.get_level_values("datetime") aligned = bl.reindex(pd.DatetimeIndex(dt.unique())).reindex(dt) mask = aligned.notna().values out = df.copy() for c in cols: vals = df[c].values res = vals.copy() res[mask] = np.asarray(vals[mask], dtype=float) - aligned[mask].values out[c] = res return out class BenchBetaResidual(processor_module.Processor): """Residualize the label against a beta-scaled benchmark move: ``r_i - b_i * r_bench``. Unlike a plain constant subtraction (see ``BenchResidual``), the name-specific rolling beta ``b_i`` makes this survive cross-sectional normalization: in up-weeks high-beta names lose rank, in down-weeks they gain — exactly the relative structure an absolute- return ranking hides. Beta is estimated from *past* data only (rolling ``window`` trading days of daily close returns of each instrument vs the benchmark, both read up to and including ``t``), so no lookahead enters the target. The benchmark leg uses the same horizon as the label expression (``Ref($close,-6)/Ref($close,-1)-1`` by default via ``horizon``/``base``, matching the yaml's 6-day label). Rows with missing beta or benchmark values keep their raw label. Requires ``$close`` to be present in the feature group (it always is for TACHandler). Stateless; accepts ``fit_start_time``/``fit_end_time`` (ignored) for uniform kwargs injection. Must be inserted BEFORE any per-date normalization processor. """ def __init__( self, benchmark="SPY", fields_group="label", lake_root=None, market="US", timeframe=None, freq="day", window=63, horizon=6, base=1, feature_field="$close", fit_start_time=None, fit_end_time=None, ): self.benchmark = benchmark self.fields_group = fields_group self.lake_root = lake_root self.market = market self.timeframe = timeframe or timeframe_for_freq(freq) self.window = int(window) self.horizon = int(horizon) self.base = int(base) self.feature_field = feature_field self.fit_start_time = fit_start_time self.fit_end_time = fit_end_time self._bench = None def _load_bench_close(self): if self._bench is not None: return self._bench cfg = LakeConfig(self.lake_root, self.market) p = cfg.bar_path(self.timeframe, self.benchmark) if not p.exists(): raise FileNotFoundError(f"BenchBetaResidual: benchmark bar file not found: {p}") df = pd.read_parquet(p) s = pd.Series(df["c"].astype(float).values, index=pd.to_datetime(df["t"])).sort_index() s.index = s.index.normalize() self._bench = s[~s.index.duplicated(keep="last")] return self._bench def fit(self, df=None): return self def __call__(self, df): bench = self._load_bench_close() # benchmark forward return over the same horizon as the label expression fwd = bench.shift(-(self.base + self.horizon - 1)) / bench.shift(-self.base) - 1 px_col = ("feature", self.feature_field) if px_col not in df.columns: raise KeyError(f"BenchBetaResidual: {self.feature_field} not found in features") px = df[px_col].unstack("instrument").sort_index() rets = px / px.shift(1) - 1 bret = bench.reindex(px.index).pct_change() # rolling beta per instrument using data <= t (no lookahead) cov = rets.rolling(self.window, min_periods=max(10, self.window // 2)).cov(bret) var = bret.rolling(self.window, min_periods=max(10, self.window // 2)).var() beta = cov.div(var, axis=0) contrib = beta.mul(fwd.reindex(px.index), axis=0) cols = list(processor_module.get_group_columns(df, self.fields_group)) out = df.copy() for c in cols: lab = df[c].unstack("instrument").reindex(px.index) resid = lab - contrib.where(contrib.notna() & lab.notna(), 0.0) new_vals = resid.stack() new_vals.index.names = df.index.names # residual where available, raw label otherwise (e.g. beta warm-up rows) out[c] = new_vals.reindex(out.index).fillna(df[c]) return out class TACHandler(DataHandlerLP): """DataHandlerLP backed by the TradeAC parquet lake. Parameters mirror ``Alpha158``: ``instruments``/``start_time``/``end_time``/``freq`` define the queried window; ``feature_fields`` selects the features (default: raw OHLCV + all common ta-lib columns found in the lake); ``label`` is a qlib expression for the target. """ def __init__( self, instruments="all", start_time=None, end_time=None, freq="day", infer_processors=DEFAULT_INFER_PROCESSORS, learn_processors=DEFAULT_LEARN_PROCESSORS, fit_start_time=None, fit_end_time=None, process_type=DataHandlerLP.PTYPE_A, filter_pipe=None, feature_fields=None, label=DEFAULT_LABEL, lake_root=None, market="US", **kwargs, ): # default the processor fit window to the queried window (like Alpha158 without a split) if fit_start_time is None: fit_start_time = start_time if fit_end_time is None: fit_end_time = end_time infer_processors = check_transform_proc(infer_processors, fit_start_time, fit_end_time) learn_processors = check_transform_proc(learn_processors, fit_start_time, fit_end_time) feature_fields = self._normalize_feature_fields(feature_fields, freq, lake_root, market) if not feature_fields: raise ValueError( "no feature fields available for the lake; set `feature_fields` explicitly " "(e.g. ['$close', '$rsi_14', '$sma_20'])" ) label_expr, label_names = self._normalize_label(label) data_loader = { "class": "QlibDataLoader", "kwargs": { "config": { "feature": (feature_fields, feature_fields), "label": (label_expr, label_names), }, "filter_pipe": filter_pipe, "freq": freq, }, } super().__init__( instruments=instruments, start_time=start_time, end_time=end_time, data_loader=data_loader, infer_processors=infer_processors, learn_processors=learn_processors, process_type=process_type, **kwargs, ) # ------------------------------------------------------------------ config @staticmethod def _normalize_feature_fields(feature_fields, freq, lake_root, market) -> List[str]: if feature_fields is None: common = get_common_feature_fields(lake_root, market, timeframe_for_freq(freq)) feature_fields = list(RAW_FEATURE_FIELDS) + ["$" + f for f in common if "$" + f not in RAW_FEATURE_FIELDS] elif isinstance(feature_fields, str): feature_fields = [f.strip() for f in feature_fields.split(",") if f.strip()] fields = [f if f.startswith("$") else "$" + f for f in feature_fields] # de-dup while preserving order seen, out = set(), [] for f in fields: if f not in seen: seen.add(f) out.append(f) return out @staticmethod def _normalize_label(label) -> Tuple[List[str], List[str]]: if isinstance(label, str): return [label], ["LABEL0"] if isinstance(label, (list, tuple)): if len(label) == 2 and isinstance(label[0], str): return [label[0]], list(label[1]) if isinstance(label[1], (list, tuple)) else [label[1]] return list(label), ["LABEL%d" % i for i in range(len(label))] raise TypeError(f"unsupported label config: {label!r}") # ------------------------------------------------------------------ utils def get_label_config(self): return DEFAULT_LABEL @staticmethod def discover_feature_fields(lake_root=None, market="US", freq="day") -> List[str]: return get_common_feature_fields(lake_root, market, timeframe_for_freq(freq)) __all__ = ["TACHandler", "DropAllNaN", "BenchResidual", "BenchBetaResidual", "get_common_feature_fields"] # Make `DropAllNaN`/`BenchResidual`/`BenchBetaResidual` resolvable by bare name from processor # configs (e.g. the default ``infer_processors`` and workflow yamls that reference them without a # ``module_path``), mirroring how qlib registers its own processors in ``qlib.data.dataset.processor``. processor_module.DropAllNaN = DropAllNaN processor_module.BenchResidual = BenchResidual processor_module.BenchBetaResidual = BenchBetaResidual