stock-embed-mamba3 / features.py
HamSession's picture
Add features.py
5003d0b verified
Raw History Blame Contribute Delete
2.53 kB
"""Build the nine model input channels from adjusted 1-minute bars.
Reference implementation shipped with the model (pandas and numpy only). The formulas,
clips and conventions are documented in preprocessing_contract.md. This file is the
executable form of that table. It was checked against the training pipeline's
implementation with zero numerical difference.
"""
from __future__ import annotations
import numpy as np
import pandas as pd
EASTERN = "America/New_York"
CHANNELS = [
"log_return",
"hl_range",
"oc_body",
"log_vol_ratio",
"session_flag",
"tod_sin",
"tod_cos",
"dow_sin",
"dow_cos",
]
def make_features(df: pd.DataFrame) -> np.ndarray:
"""Nine channels for one security.
Args:
df: one security sorted by `timestamp_eastern`, with columns `open_adj`,
`high_adj`, `low_adj`, `close_adj`, `volume_adj`, `timestamp_eastern`
(tz-aware; any zone is converted to Eastern, naive is taken as Eastern).
Returns:
(len(df), 9) float32 array in CHANNELS order.
"""
ts = df["timestamp_eastern"]
if not pd.api.types.is_datetime64_any_dtype(ts):
ts = pd.to_datetime(ts)
ts = ts.dt.tz_localize(EASTERN) if ts.dt.tz is None else ts.dt.tz_convert(EASTERN)
date = ts.dt.date
prev_close = df["close_adj"].shift(1)
denom = prev_close.fillna(df["close_adj"])
log_return = np.log(df["close_adj"] / prev_close).fillna(0.0).clip(-0.2, 0.2)
hl_range = ((df["high_adj"] - df["low_adj"]) / denom).clip(0.0, 0.2)
oc_body = ((df["close_adj"] - df["open_adj"]) / denom).clip(-0.2, 0.2)
daily_mean_vol = df.groupby(date)["volume_adj"].mean()
rolling = daily_mean_vol.rolling(20, min_periods=1).mean().shift(1)
rolling_per_bar = date.map(rolling)
bad_mean = rolling_per_bar.isna() | (rolling_per_bar == 0)
ratio = (df["volume_adj"] / rolling_per_bar).where(~bad_mean & (df["volume_adj"] > 0), 1.0)
log_vol_ratio = np.clip(np.where(df["volume_adj"] == 0, -10.0, np.log(ratio)), -10.0, 10.0)
session_flag = (date != date.shift(1)).astype(float)
tod = (ts.dt.hour + ts.dt.minute / 60.0) * (2 * np.pi / 24.0)
dow = ts.dt.weekday * (2 * np.pi / 5.0)
feats = np.stack(
[
log_return,
hl_range,
oc_body,
log_vol_ratio,
session_flag,
np.sin(tod),
np.cos(tod),
np.sin(dow),
np.cos(dow),
],
axis=1,
)
return feats.astype(np.float32)