ndsideload/stk_predt

1"""2Data Manager — fetch, clean, and serve the historical stock dataset.3 4Reads the published Google Sheet CSV, handles type coercion, filters to5a liquid universe, and provides date-indexed panels for downstream models.6"""7from __future__ import annotations8 9import io10import logging11from typing import Optional12 13import numpy as np14import pandas as pd15import requests16 17import config18 19logger = logging.getLogger(__name__)20 21 22def fetch_csv(url: Optional[str] = None) -> pd.DataFrame:23    """Download the published CSV and return a cleaned DataFrame."""24    url = url or config.SOURCE_CSV_URL25    logger.info("Fetching CSV from %s", url[:80])26    resp = requests.get(url, timeout=60)27    resp.raise_for_status()28    df = pd.read_csv(io.StringIO(resp.text))29    return _clean(df)30 31 32def load_csv(path: str) -> pd.DataFrame:33    """Load from a local CSV file (for backtesting / dev)."""34    df = pd.read_csv(path)35    return _clean(df)36 37 38def _clean(df: pd.DataFrame) -> pd.DataFrame:39    """Type coercion and basic sanity checks."""40    df["Log Date"] = pd.to_datetime(df["Log Date"])41    # Force numeric on Marketcap/volume (may contain #DIV/0!)42    df["Marketcap/volume"] = pd.to_numeric(df["Marketcap/volume"], errors="coerce")43    # Coerce all feature columns44    for col in config.FEATURE_COLS:45        if col in df.columns:46            df[col] = pd.to_numeric(df[col], errors="coerce")47    df.sort_values(["Log Date", "Stock Symbol"], inplace=True)48    df.reset_index(drop=True, inplace=True)49    logger.info(50        "Loaded %d rows | %d stocks | %d dates",51        len(df), df["Stock Symbol"].nunique(), df["Log Date"].nunique(),52    )53    return df54 55 56def filter_liquid(df: pd.DataFrame) -> pd.DataFrame:57    """Keep only stocks meeting minimum market-cap and volume thresholds.58    When MIN_MCAP and MIN_VOLUME are both 0, returns all stocks (no filtering)."""59    if config.MIN_MCAP <= 0 and config.MIN_VOLUME <= 0:60        logger.info("No liquidity filter — using all %d stocks", df["Stock Symbol"].nunique())61        return df.copy()62 63    latest = df["Log Date"].max()64    latest_snap = df[df["Log Date"] == latest]65    valid_symbols = latest_snap[66        (latest_snap["Market Cap"].fillna(0) >= config.MIN_MCAP)67        & (latest_snap["Volume"].fillna(0) >= config.MIN_VOLUME)68    ]["Stock Symbol"].unique()69    filtered = df[df["Stock Symbol"].isin(valid_symbols)].copy()70    logger.info("Liquid universe: %d stocks (from %d)", len(valid_symbols), df["Stock Symbol"].nunique())71    return filtered72 73 74def build_panel(df: pd.DataFrame) -> dict[str, pd.DataFrame]:75    """76    Return a dict mapping stock symbol → DataFrame indexed by date.77    Each stock-df has feature columns + 'target_direction' + 'target_return'.78    NOTE: Drops the latest row (no next-day target available) — used for TRAINING.79    """80    panel: dict[str, pd.DataFrame] = {}81    for sym, grp in df.groupby("Stock Symbol"):82        g = grp.set_index("Log Date").sort_index()83        g["next_price"] = g["Today Price"].shift(-1)84        g["target_return"] = (g["next_price"] / g["Today Price"] - 1) * 10085        g["target_direction"] = (g["target_return"] > 0).astype(int)86        g.dropna(subset=["target_return"], inplace=True)87        if len(g) >= 2:88            panel[sym] = g89    logger.info("Built panel for %d stocks with targets", len(panel))90    return panel91 92 93def build_full_panel(df: pd.DataFrame) -> dict[str, pd.DataFrame]:94    """95    Like build_panel but KEEPS the latest row (with NaN target).96    Used for building prediction sequences — the LSTM/Transformer need97    the most recent day's features even though there's no target for it.98    """99    panel: dict[str, pd.DataFrame] = {}100    for sym, grp in df.groupby("Stock Symbol"):101        g = grp.set_index("Log Date").sort_index()102        g["next_price"] = g["Today Price"].shift(-1)103        g["target_return"] = (g["next_price"] / g["Today Price"] - 1) * 100104        g["target_direction"] = (g["target_return"] > 0).astype(int)105        # Fill NaN targets on last row with 0 (placeholder — won't be used for training)106        g["target_return"] = g["target_return"].fillna(0)107        g["target_direction"] = g["target_direction"].fillna(0)108        if len(g) >= 2:109            panel[sym] = g110    logger.info("Built full panel for %d stocks (includes latest date)", len(panel))111    return panel112 113 114def get_latest_snapshot(df: pd.DataFrame) -> pd.DataFrame:115    """Return only the most recent date's data (for live prediction)."""116    latest = df["Log Date"].max()117    return df[df["Log Date"] == latest].copy().set_index("Stock Symbol")118 119 120def top_n_by_mcap(df: pd.DataFrame, n: int = config.HEAVY_MODEL_TOP_N) -> list[str]:121    """Return top-N symbols by market cap on the latest date. If n >= total stocks, returns all."""122    snap = get_latest_snapshot(df)123    if n >= len(snap):124        return snap.index.tolist()125    return snap.nlargest(n, "Market Cap").index.tolist()126