Skip to content

Commit 6546912

Browse files
casperngeencasperngeen
committed
Add baseline model for Kalshi event prediction
1 parent 7145f9a commit 6546912

4 files changed

Lines changed: 227 additions & 0 deletions

File tree

‎kalshi/__init__.py‎

Whitespace-only changes.

‎kalshi/baseline.py‎

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
"""Kalshi Economic Events — Implied Mean Surprise Direction Baseline.
2+
3+
Recovers implied means from Kalshi threshold-based event markets
4+
(CPI, Core CPI, FED level/decision, Payrolls, Unemployment, GDP)
5+
via the stg_infra events submodule and trains a logistic regression
6+
baseline to predict surprise direction:
7+
8+
label=1 → actual print > implied mean (upside surprise)
9+
label=0 → actual print < implied mean (downside surprise)
10+
11+
Features come directly from the PDF-derived statistics returned by the
12+
stg_infra event series functions: implied_mean, implied_std, implied_skew,
13+
implied_kurtosis, implied_entropy, implied_median, n_submarkets.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
import logging
19+
import sys
20+
from pathlib import Path
21+
22+
from dotenv import load_dotenv
23+
load_dotenv(Path(__file__).parent / ".env")
24+
25+
# stg_infra uses two import roots:
26+
# - PortfolioBench/stg → makes `stg_infra.*` findable (namespace pkg)
27+
# - PortfolioBench/stg/stg_infra → makes `stg.*` findable (installed pkg root)
28+
_STG_ROOT = Path(__file__).parent.parent / "stg"
29+
sys.path.insert(0, str(_STG_ROOT))
30+
sys.path.insert(0, str(_STG_ROOT / "stg_infra"))
31+
32+
import numpy as np
33+
import polars as pl
34+
from sklearn.linear_model import LogisticRegression
35+
from sklearn.metrics import classification_report, accuracy_score
36+
from sklearn.preprocessing import StandardScaler
37+
38+
from stg.events.cpi import compute_cpi_mom_series, compute_cpi_yoy_series
39+
from stg.events.core_cpi import compute_core_cpi_mom_series, compute_core_cpi_yoy_series
40+
from stg.events.fed import compute_fed_level_series, compute_fed_decision_series
41+
from stg.events.payrolls import compute_payrolls_series
42+
from stg.events.unemployment import compute_unemployment_series
43+
from stg.events.gdp import compute_gdp_series
44+
45+
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
46+
log = logging.getLogger(__name__)
47+
48+
DATA_DIR = Path("user_data/data/kalshi/kalshi")
49+
TRAIN_RATIO = 0.7
50+
51+
SERIES_ORDER = [
52+
"cpi_mom", "cpi_yoy",
53+
"core_cpi_mom", "core_cpi_yoy",
54+
"fed_level", "decision_bps",
55+
"payrolls",
56+
"unemployment",
57+
"gdp",
58+
]
59+
60+
# ---------------------------------------------------------------------------
61+
# 1. Load data
62+
# ---------------------------------------------------------------------------
63+
log.info("Loading markets...")
64+
markets = pl.read_parquet(str(DATA_DIR / "markets/*.parquet"))
65+
66+
log.info("Loading trades...")
67+
trades = pl.read_parquet(str(DATA_DIR / "trades/*.parquet"))
68+
69+
log.info("Markets: %d rows | Trades: %d rows", len(markets), len(trades))
70+
71+
# ---------------------------------------------------------------------------
72+
# 2. Compute implied mean series for every supported event type
73+
# ---------------------------------------------------------------------------
74+
log.info("Computing implied mean series...")
75+
76+
_series_fns = [
77+
compute_cpi_mom_series,
78+
compute_cpi_yoy_series,
79+
compute_core_cpi_mom_series,
80+
compute_core_cpi_yoy_series,
81+
compute_fed_level_series,
82+
compute_fed_decision_series,
83+
compute_payrolls_series,
84+
compute_unemployment_series,
85+
compute_gdp_series,
86+
]
87+
88+
frames = []
89+
for fn in _series_fns:
90+
df = fn(markets, trades)
91+
if not df.is_empty():
92+
frames.append(df)
93+
log.info(" %-25s %d rows %d events",
94+
fn.__name__, len(df), df["event_ticker"].n_unique())
95+
96+
if not frames:
97+
sys.exit("No implied mean data found.")
98+
99+
implied = pl.concat(frames, how="diagonal").sort(["event_ticker", "date"])
100+
log.info("Total implied-mean rows: %d across %d events",
101+
len(implied), implied["event_ticker"].n_unique())
102+
103+
# ---------------------------------------------------------------------------
104+
# 3. Label each observation
105+
# ---------------------------------------------------------------------------
106+
# Keep only rows where resolved_value is known and implied_mean is not null
107+
# and there is a clear surprise (skip exact ties)
108+
labeled = (
109+
implied
110+
.filter(
111+
pl.col("resolved_value").is_not_null()
112+
& pl.col("implied_mean").is_not_null()
113+
& (pl.col("resolved_value") != pl.col("implied_mean"))
114+
)
115+
.with_columns(
116+
pl.when(pl.col("resolved_value") > pl.col("implied_mean"))
117+
.then(pl.lit(1))
118+
.otherwise(pl.lit(0))
119+
.alias("label")
120+
)
121+
)
122+
123+
log.info("Labeled observations: %d (up=%d, down=%d)",
124+
len(labeled),
125+
(labeled["label"] == 1).sum(),
126+
(labeled["label"] == 0).sum())
127+
128+
if labeled.is_empty():
129+
sys.exit("No labeled observations after filtering.")
130+
131+
# ---------------------------------------------------------------------------
132+
# 4. Chronological train / test split (split on date, not shuffle)
133+
# ---------------------------------------------------------------------------
134+
all_dates = labeled["date"].sort()
135+
split_idx = int(len(all_dates) * TRAIN_RATIO)
136+
split_date = all_dates[split_idx]
137+
138+
train_df = labeled.filter(pl.col("date") < split_date)
139+
test_df = labeled.filter(pl.col("date") >= split_date)
140+
141+
log.info("Train: %d obs (%s → %s)", len(train_df),
142+
train_df["date"].min(), train_df["date"].max())
143+
log.info("Test: %d obs (%s → %s)", len(test_df),
144+
test_df["date"].min(), test_df["date"].max())
145+
146+
# ---------------------------------------------------------------------------
147+
# 5. Assemble feature matrix
148+
# ---------------------------------------------------------------------------
149+
FEATURE_COLS = [
150+
"implied_mean",
151+
"implied_std",
152+
"implied_skew",
153+
"implied_kurtosis",
154+
"implied_entropy",
155+
"implied_median",
156+
"n_submarkets",
157+
]
158+
159+
160+
def to_xy(df: pl.DataFrame) -> tuple[np.ndarray, np.ndarray]:
161+
X = df.select(FEATURE_COLS).to_numpy().astype(np.float64)
162+
y = df["label"].to_numpy().astype(np.int64)
163+
# Drop rows with any NaN feature (lag windows produce NaNs at start of series)
164+
mask = ~np.isnan(X).any(axis=1)
165+
return X[mask], y[mask]
166+
167+
168+
X_train, y_train = to_xy(train_df)
169+
X_test, y_test = to_xy(test_df)
170+
171+
log.info("Train samples after NaN drop: %d (up=%d, down=%d)",
172+
len(y_train), (y_train == 1).sum(), (y_train == 0).sum())
173+
log.info("Test samples after NaN drop: %d (up=%d, down=%d)",
174+
len(y_test), (y_test == 1).sum(), (y_test == 0).sum())
175+
176+
if len(np.unique(y_train)) < 2:
177+
sys.exit("Training set has only one class — cannot train classifier.")
178+
179+
# ---------------------------------------------------------------------------
180+
# 6. Train logistic regression baseline
181+
# ---------------------------------------------------------------------------
182+
scaler = StandardScaler()
183+
X_train_s = scaler.fit_transform(X_train)
184+
X_test_s = scaler.transform(X_test)
185+
186+
model = LogisticRegression(max_iter=1000, class_weight="balanced", random_state=42)
187+
model.fit(X_train_s, y_train)
188+
189+
y_pred = model.predict(X_test_s)
190+
191+
# ---------------------------------------------------------------------------
192+
# 7. Report
193+
# ---------------------------------------------------------------------------
194+
print("\n" + "=" * 60)
195+
print("Implied Mean Surprise Direction Baseline (Logistic Regression)")
196+
print("=" * 60)
197+
print(f"Test accuracy: {accuracy_score(y_test, y_pred):.4f}")
198+
print()
199+
print(classification_report(y_test, y_pred, target_names=["downside", "upside"]))
200+
201+
coefs = model.coef_[0]
202+
print("Feature coefficients (sorted by magnitude):")
203+
for name, coef in sorted(zip(FEATURE_COLS, coefs), key=lambda x: abs(x[1]), reverse=True):
204+
print(f" {name:20s} {coef:+.4f}")
205+
206+
# Per-series breakdown
207+
print("\nPer-series accuracy:")
208+
for stype in SERIES_ORDER:
209+
sub = test_df.filter(pl.col("series_type") == stype)
210+
if sub.is_empty():
211+
continue
212+
Xs, ys = to_xy(sub)
213+
if len(Xs) == 0:
214+
continue
215+
Xs_s = scaler.transform(Xs)
216+
preds = model.predict(Xs_s)
217+
acc = accuracy_score(ys, preds)
218+
print(f" {stype:20s} n={len(ys):4d} acc={acc:.4f}")

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ dependencies = [
6666
"aiohttp",
6767
"cryptography",
6868
"sdnotify",
69+
"python-dotenv",
6970
"python-dateutil",
7071
"pytz",
7172
"packaging",

‎pyrightconfig.json‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"venvPath": ".",
3+
"venv": ".venv",
4+
"extraPaths": [
5+
"stg/stg_infra",
6+
"stg"
7+
]
8+
}

0 commit comments

Comments
 (0)