From a2f67ae9ed6b76400fd4332f28ce5de4003f0433 Mon Sep 17 00:00:00 2001 From: Tarandeep Singh Juneja Date: Sat, 15 Aug 2026 16:03:30 +0530 Subject: [PATCH] Add sample-path quantiles and a random-walk walk-forward eval for #387. Zero-shot Kronos-small was being judged as a single-name directional signal, while predict() averaged sample paths and hid the real forecast cone. Keep the mean path for compatibility, return quantiles from forecast(), and compare MAE/hit-rate/coverage against a last-price random walk. Co-authored-by: Cursor --- README.md | 12 ++ examples/walk_forward_eval.py | 239 ++++++++++++++++++++++++++++++++ model/__init__.py | 2 +- model/kronos.py | 111 +++++++++++++-- tests/test_walk_forward_eval.py | 92 ++++++++++++ 5 files changed, 444 insertions(+), 12 deletions(-) create mode 100644 examples/walk_forward_eval.py create mode 100644 tests/test_walk_forward_eval.py diff --git a/README.md b/README.md index faffa1578..39cc1de62 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,18 @@ Running this script will generate a plot comparing the ground truth data against Additionally, we provide a script that makes predictions without Volume and Amount data, which can be found in [`examples/prediction_wo_vol_example.py`](examples/prediction_wo_vol_example.py). +### Evaluating forecasts (walk-forward, issue [#387](https://github.com/shiyu-coder/Kronos/issues/387)) + +Zero-shot Kronos-small is a **generative K-line prior**, not a finished single-contract trading system. The paper reports IC / RankIC (series correlation and cross-sectional ranking). A walk-forward long/short on one futures series can easily show **no directional edge and worse MAE than a last-price random walk** — especially if you plot the high/low of the *averaged* `predict()` path as a confidence cone. That cone is too narrow because `predict()` averages `sample_count` trajectories internally. + +Use `KronosPredictor.forecast(...)` to keep sample paths and 10/90% quantiles, and compare against a random-walk baseline on the same origins: + +```shell +python examples/walk_forward_eval.py --data ./data/XSHG_5min_600977.csv --pred-len 12 --windows 8 --sample-count 8 +``` + +If that script reports no edge on your market, finetune on that instrument/frequency (see below) rather than only raising `T`. + ## 🔧 Finetuning on Your Own Data (A-Share Market Example) diff --git a/examples/walk_forward_eval.py b/examples/walk_forward_eval.py new file mode 100644 index 000000000..7b1905c1e --- /dev/null +++ b/examples/walk_forward_eval.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +"""Walk-forward evaluation of zero-shot Kronos vs a last-price random walk. + +This addresses https://github.com/shiyu-coder/Kronos/issues/387: + +* Single-name directional PnL is not the paper protocol. The paper reports IC / + RankIC of predicted vs actual series (and cross-sectional ranking after + finetuning), not "long/short one futures contract vs a coin flip". +* `predict()` averages sample paths. Plotting high/low of that mean path makes + uncertainty cones look too narrow. Use `forecast()` quantiles across samples. +* Always compare MAE and hit-rate against a last-price random walk on the + *same* origins before claiming an edge. + +Example: + python examples/walk_forward_eval.py \\ + --data data/XSHG_5min_600977.csv \\ + --lookback 400 --pred-len 12 --stride 48 --windows 8 --sample-count 8 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from model import Kronos, KronosPredictor, KronosTokenizer + + +def directional_hit_rate(pred_ret: np.ndarray, actual_ret: np.ndarray) -> float: + pred_ret = np.asarray(pred_ret, dtype=np.float64) + actual_ret = np.asarray(actual_ret, dtype=np.float64) + mask = (pred_ret != 0) & (actual_ret != 0) & np.isfinite(pred_ret) & np.isfinite(actual_ret) + if mask.sum() == 0: + return float("nan") + return float(np.mean(np.sign(pred_ret[mask]) == np.sign(actual_ret[mask]))) + + +def pearson_ic(pred_ret: np.ndarray, actual_ret: np.ndarray) -> float: + pred_ret = np.asarray(pred_ret, dtype=np.float64) + actual_ret = np.asarray(actual_ret, dtype=np.float64) + if pred_ret.size < 3: + return float("nan") + if np.std(pred_ret) == 0 or np.std(actual_ret) == 0: + return float("nan") + return float(np.corrcoef(pred_ret, actual_ret)[0, 1]) + + +def cone_coverage(actual: np.ndarray, q_low: np.ndarray, q_high: np.ndarray) -> float: + actual = np.asarray(actual, dtype=np.float64) + return float(np.mean((actual >= q_low) & (actual <= q_high))) + + +def naive_mean_path_cone_width(samples: np.ndarray) -> float: + """Width you get if you average paths first, then read high-low of the mean bar.""" + mean_path = samples.mean(axis=0) + return float(np.mean(mean_path[:, 1] - mean_path[:, 2])) + + +def sample_quantile_cone_width(samples: np.ndarray, q_low=0.1, q_high=0.9, close_idx=3) -> float: + """Width of the close distribution across samples (the actual forecast cone).""" + low = np.quantile(samples[:, :, close_idx], q_low, axis=0) + high = np.quantile(samples[:, :, close_idx], q_high, axis=0) + return float(np.mean(high - low)) + + +def summarize_windows(rows: list[dict]) -> dict: + pred = np.array([r["kronos_return"] for r in rows], dtype=np.float64) + actual = np.array([r["actual_return"] for r in rows], dtype=np.float64) + rw = np.array([r["rw_return"] for r in rows], dtype=np.float64) + mae_kronos = float(np.mean([r["kronos_mae"] for r in rows])) + mae_rw = float(np.mean([r["rw_mae"] for r in rows])) + coverage = float(np.mean([r["coverage"] for r in rows])) + naive_width = float(np.mean([r["naive_cone_width"] for r in rows])) + sample_width = float(np.mean([r["sample_cone_width"] for r in rows])) + return { + "n_windows": len(rows), + "kronos_mae": mae_kronos, + "random_walk_mae": mae_rw, + "mae_ratio_vs_rw": mae_kronos / mae_rw if mae_rw else float("nan"), + "kronos_hit_rate": directional_hit_rate(pred, actual), + "random_walk_hit_rate": directional_hit_rate(rw, actual), + "kronos_ic": pearson_ic(pred, actual), + "q10_q90_coverage": coverage, + "naive_mean_path_cone_width": naive_width, + "sample_quantile_cone_width": sample_width, + "cone_width_ratio_naive_over_sample": naive_width / sample_width if sample_width else float("nan"), + } + + +def interpret(summary: dict) -> list[str]: + lines = [] + ratio = summary["mae_ratio_vs_rw"] + if np.isfinite(ratio) and ratio >= 1.0: + lines.append( + f"No point-forecast edge: Kronos MAE is {ratio:.3f}x the last-price random walk." + ) + elif np.isfinite(ratio): + lines.append( + f"Kronos MAE is {ratio:.3f}x the last-price random walk (lower is better)." + ) + hit = summary["kronos_hit_rate"] + if np.isfinite(hit) and hit <= 0.5: + lines.append( + f"No directional edge: hit rate {hit:.1%} is not above a coin flip." + ) + cov = summary["q10_q90_coverage"] + if np.isfinite(cov) and cov < 0.7: + lines.append( + f"Cones still look tight: {cov:.1%} of actuals fell inside the 10-90% sample band " + "(80% expected if the band is calibrated)." + ) + width_ratio = summary["cone_width_ratio_naive_over_sample"] + if np.isfinite(width_ratio) and width_ratio < 0.5: + lines.append( + f"The mean-path high/low cone is {width_ratio:.2f}x the sample-quantile cone. " + "Do not plot high/low of predict() as uncertainty." + ) + ic = summary["kronos_ic"] + if np.isfinite(ic): + lines.append( + f"Single-series IC of horizon returns is {ic:.4f}. The paper's RankIC is a " + "cross-sectional ranking metric after (usually) finetuning, not this test." + ) + lines.append( + "Zero-shot Kronos-small is a generative K-line prior, not a finished trading signal. " + "If this walk-forward shows no edge on your futures, finetune on that market " + "or use cross-sectional ranking — do not expect T=1.5 alone to create alpha." + ) + return lines + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Walk-forward Kronos vs random-walk evaluation") + parser.add_argument("--data", type=Path, default=ROOT / "data" / "XSHG_5min_600977.csv") + parser.add_argument("--lookback", type=int, default=400) + parser.add_argument("--pred-len", type=int, default=12) + parser.add_argument("--stride", type=int, default=48) + parser.add_argument("--windows", type=int, default=6) + parser.add_argument("--sample-count", type=int, default=8) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--top-p", type=float, default=0.9) + parser.add_argument("--tokenizer", default="NeoQuasar/Kronos-Tokenizer-base") + parser.add_argument("--model", default="NeoQuasar/Kronos-small") + parser.add_argument("--output", type=Path, default=None) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if not args.data.exists(): + print(f"Data file not found: {args.data}") + return 1 + + df = pd.read_csv(args.data) + time_col = "timestamps" if "timestamps" in df.columns else "timestamp" + df[time_col] = pd.to_datetime(df[time_col]) + cols = [c for c in ["open", "high", "low", "close", "volume", "amount"] if c in df.columns] + + tokenizer = KronosTokenizer.from_pretrained(args.tokenizer) + model = Kronos.from_pretrained(args.model) + predictor = KronosPredictor(model, tokenizer, max_context=512) + print(f"Device: {predictor.device}") + + max_origin = len(df) - args.pred_len + origins = list(range(args.lookback, max_origin, args.stride))[: args.windows] + if not origins: + print("Not enough rows for the requested lookback/pred_len.") + return 1 + + rows = [] + close_idx = cols.index("close") + for origin in origins: + context = df.iloc[origin - args.lookback:origin] + future = df.iloc[origin:origin + args.pred_len] + x_df = context[cols].reset_index(drop=True) + x_ts = context[time_col].reset_index(drop=True) + y_ts = future[time_col].reset_index(drop=True) + + result = predictor.forecast( + df=x_df, + x_timestamp=x_ts, + y_timestamp=y_ts, + pred_len=args.pred_len, + T=args.temperature, + top_p=args.top_p, + sample_count=args.sample_count, + verbose=False, + ) + last_close = float(context["close"].iloc[-1]) + actual_close = future["close"].to_numpy(dtype=np.float64) + mean_close = result["mean"]["close"].to_numpy(dtype=np.float64) + q10 = result["quantiles"]["q10"]["close"].to_numpy(dtype=np.float64) + q90 = result["quantiles"]["q90"]["close"].to_numpy(dtype=np.float64) + rw_path = np.repeat(last_close, args.pred_len) + + kronos_end = float(mean_close[-1]) + actual_end = float(actual_close[-1]) + rows.append({ + "origin": str(context[time_col].iloc[-1]), + "last_close": last_close, + "kronos_return": kronos_end / last_close - 1.0, + "actual_return": actual_end / last_close - 1.0, + "rw_return": 0.0, + "kronos_mae": float(np.mean(np.abs(mean_close - actual_close))), + "rw_mae": float(np.mean(np.abs(rw_path - actual_close))), + "coverage": cone_coverage(actual_close, q10, q90), + "naive_cone_width": naive_mean_path_cone_width(result["samples"]), + "sample_cone_width": sample_quantile_cone_width(result["samples"], close_idx=close_idx), + }) + print( + f"{rows[-1]['origin']} kronos_ret={rows[-1]['kronos_return']:+.4f} " + f"actual_ret={rows[-1]['actual_return']:+.4f} " + f"mae={rows[-1]['kronos_mae']:.4f} vs rw={rows[-1]['rw_mae']:.4f} " + f"cover={rows[-1]['coverage']:.0%}" + ) + + summary = summarize_windows(rows) + print("\n=== Walk-forward summary ===") + print(json.dumps(summary, indent=2)) + print("\n=== Interpretation ===") + for line in interpret(summary): + print(f"- {line}") + + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps({"windows": rows, "summary": summary}, indent=2)) + print(f"\nWrote {args.output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/model/__init__.py b/model/__init__.py index 718d07a21..5d2e61f6d 100644 --- a/model/__init__.py +++ b/model/__init__.py @@ -1,4 +1,4 @@ -from .kronos import KronosTokenizer, Kronos, KronosPredictor +from .kronos import KronosTokenizer, Kronos, KronosPredictor, aggregate_sample_paths model_dict = { 'kronos_tokenizer': KronosTokenizer, diff --git a/model/kronos.py b/model/kronos.py index ce4494ee0..17d2943a6 100644 --- a/model/kronos.py +++ b/model/kronos.py @@ -386,7 +386,28 @@ def sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_l return x -def auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False): +def aggregate_sample_paths(samples, quantile_levels=(0.1, 0.5, 0.9)): + """Reduce (n_samples, pred_len, n_features) paths into mean/std/quantiles. + + Averaging paths first and then reading high/low of that mean path makes + uncertainty cones look far too narrow. Quantiles must be taken across samples. + """ + samples = np.asarray(samples, dtype=np.float64) + if samples.ndim != 3: + raise ValueError(f"samples must have shape (n_samples, pred_len, n_features), got {samples.shape}") + n_samples = samples.shape[0] + summary = { + "mean": samples.mean(axis=0), + "std": samples.std(axis=0, ddof=1) if n_samples > 1 else np.zeros(samples.shape[1:], dtype=np.float64), + "n_samples": n_samples, + } + for q in quantile_levels: + key = f"q{int(round(float(q) * 100))}" + summary[key] = np.quantile(samples, q, axis=0) + return summary + + +def auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False, return_samples=False): with torch.no_grad(): x = torch.clip(x, -clip, clip) @@ -464,9 +485,9 @@ def auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context z = tokenizer.decode(input_tokens, half=True) z = z.reshape(-1, sample_count, z.size(1), z.size(2)) preds = z.cpu().numpy() - preds = np.mean(preds, axis=1) - - return preds + if return_samples: + return preds + return np.mean(preds, axis=1) def calc_time_stamps(x_timestamp): @@ -502,19 +523,25 @@ def __init__(self, model, tokenizer, device=None, max_context=512, clip=5): self.device = device - self.tokenizer = self.tokenizer.to(self.device) - self.model = self.model.to(self.device) + # Dropout must be off for inference. Leaving the module in train mode + # noisily shrinks sampled cones and can make point forecasts worse than + # a last-price random walk. + self.tokenizer = self.tokenizer.to(self.device).eval() + self.model = self.model.to(self.device).eval() - def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose): + def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, return_samples=False): x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device) x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device) y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device) - preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len, - self.clip, T, top_k, top_p, sample_count, verbose) - preds = preds[:, -pred_len:, :] - return preds + preds = auto_regressive_inference( + self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len, + self.clip, T, top_k, top_p, sample_count, verbose, return_samples=return_samples, + ) + if return_samples: + return preds[:, :, -pred_len:, :] + return preds[:, -pred_len:, :] def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True): @@ -558,6 +585,68 @@ def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p= pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp) return pred_df + def forecast(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=16, quantile_levels=(0.1, 0.5, 0.9), verbose=True): + """Probabilistic forecast that keeps sample paths for uncertainty cones. + + `predict()` averages `sample_count` paths internally, so high/low of that + single mean path is not a confidence interval. Use this method when you + need walk-forward evaluation or calibrated cones. + + Returns: + dict with: + mean (DataFrame): ensemble-mean path (same as `predict`) + quantiles (dict[str, DataFrame]): e.g. q10/q50/q90 across samples + samples (ndarray): shape (sample_count, pred_len, n_features) + last_close (float): last observed close, for return signals + columns (list[str]): feature names + """ + if sample_count < 1: + raise ValueError("sample_count must be >= 1") + + if not isinstance(df, pd.DataFrame): + raise ValueError("Input must be a pandas DataFrame.") + if not all(col in df.columns for col in self.price_cols): + raise ValueError(f"Price columns {self.price_cols} not found in DataFrame.") + + df = df.copy() + if self.vol_col not in df.columns: + df[self.vol_col] = 0.0 + df[self.amt_vol] = 0.0 + if self.amt_vol not in df.columns and self.vol_col in df.columns: + df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1) + if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any(): + raise ValueError("Input DataFrame contains NaN values in price or volume columns.") + + x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32) + x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0) + x_norm = np.clip((x - x_mean) / (x_std + 1e-5), -self.clip, self.clip) + x_stamp = calc_time_stamps(x_timestamp).values.astype(np.float32) + y_stamp = calc_time_stamps(y_timestamp).values.astype(np.float32) + + samples = self.generate( + x_norm[np.newaxis, :], + x_stamp[np.newaxis, :], + y_stamp[np.newaxis, :], + pred_len, T, top_k, top_p, sample_count, verbose, + return_samples=True, + )[0] + samples = samples * (x_std + 1e-5) + x_mean + + summary = aggregate_sample_paths(samples, quantile_levels=quantile_levels) + columns = self.price_cols + [self.vol_col, self.amt_vol] + mean_df = pd.DataFrame(summary["mean"], columns=columns, index=y_timestamp) + quantile_dfs = {} + for q in quantile_levels: + key = f"q{int(round(float(q) * 100))}" + quantile_dfs[key] = pd.DataFrame(summary[key], columns=columns, index=y_timestamp) + + return { + "mean": mean_df, + "quantiles": quantile_dfs, + "samples": samples, + "last_close": float(df["close"].iloc[-1]), + "columns": columns, + } def predict_batch(self, df_list, x_timestamp_list, y_timestamp_list, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True): """ diff --git a/tests/test_walk_forward_eval.py b/tests/test_walk_forward_eval.py new file mode 100644 index 000000000..3562df07f --- /dev/null +++ b/tests/test_walk_forward_eval.py @@ -0,0 +1,92 @@ +import sys +from pathlib import Path + +import numpy as np +import pytest + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "examples")) + +from walk_forward_eval import ( # noqa: E402 + cone_coverage, + directional_hit_rate, + interpret, + naive_mean_path_cone_width, + pearson_ic, + sample_quantile_cone_width, + summarize_windows, +) +from model.kronos import aggregate_sample_paths + + +def test_aggregate_sample_paths_quantiles_are_across_samples_not_mean_bar(): + # Three paths that fan out on close (feature 3). Averaging first collapses them. + samples = np.zeros((3, 4, 6), dtype=np.float64) + samples[0, :, 3] = 10.0 + samples[1, :, 3] = 12.0 + samples[2, :, 3] = 14.0 + samples[:, :, 1] = samples[:, :, 3] # high + samples[:, :, 2] = samples[:, :, 3] # low + + summary = aggregate_sample_paths(samples, quantile_levels=(0.1, 0.5, 0.9)) + np.testing.assert_allclose(summary["mean"][:, 3], 12.0) + assert summary["q10"][0, 3] < summary["q50"][0, 3] < summary["q90"][0, 3] + assert summary["q90"][0, 3] - summary["q10"][0, 3] > 2.0 + + naive_width = naive_mean_path_cone_width(samples) + sample_width = sample_quantile_cone_width(samples) + assert naive_width == pytest.approx(0.0) + assert sample_width > 2.0 + + +def test_directional_hit_rate_and_ic(): + pred = np.array([0.01, -0.02, 0.03, -0.01]) + actual = np.array([0.02, -0.01, -0.04, -0.02]) + assert directional_hit_rate(pred, actual) == pytest.approx(0.75) + assert pearson_ic(pred, pred) == pytest.approx(1.0) + + +def test_cone_coverage_and_summary_flags_no_edge(): + actual = np.array([10.0, 10.2, 9.8]) + assert cone_coverage(actual, np.array([9.5, 9.5, 9.5]), np.array([10.5, 10.5, 10.5])) == 1.0 + assert cone_coverage(actual, np.array([11.0, 11.0, 11.0]), np.array([12.0, 12.0, 12.0])) == 0.0 + + rows = [ + { + "kronos_return": 0.01, + "actual_return": -0.02, + "rw_return": 0.0, + "kronos_mae": 0.4, + "rw_mae": 0.2, + "coverage": 0.3, + "naive_cone_width": 0.01, + "sample_cone_width": 0.2, + }, + { + "kronos_return": -0.01, + "actual_return": 0.02, + "rw_return": 0.0, + "kronos_mae": 0.5, + "rw_mae": 0.25, + "coverage": 0.4, + "naive_cone_width": 0.01, + "sample_cone_width": 0.2, + }, + { + "kronos_return": 0.02, + "actual_return": -0.01, + "rw_return": 0.0, + "kronos_mae": 0.3, + "rw_mae": 0.15, + "coverage": 0.2, + "naive_cone_width": 0.01, + "sample_cone_width": 0.2, + }, + ] + summary = summarize_windows(rows) + assert summary["mae_ratio_vs_rw"] == pytest.approx(2.0) + notes = " ".join(interpret(summary)) + assert "No point-forecast edge" in notes + assert "No directional edge" in notes + assert "mean-path high/low" in notes + assert "RankIC" in notes