Update app.py

app.py

@@ -3,22 +3,24 @@ import math
 import tempfile
 from dataclasses import dataclass
 from functools import lru_cache
+from typing import List, Tuple
 
 import gradio as gr
+
+# Headless plotting for HF Spaces
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
+
 import numpy as np
 import pandas as pd
+import requests
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 import lightning.pytorch as pl
 from torch.utils.data import DataLoader, TensorDataset
 
-from pandas_datareader import data as pdr
-
 
 DISCLAIMER = """
 **Disclaimer (Educational Use Only):**  
@@ -29,11 +31,11 @@ Markets are risky; consult a qualified professional for investment guidance.
 
 
 # -----------------------------
-# Data & features
+# Feature engineering
 # -----------------------------
 @dataclass
 class FeatureSpec:
-    lookback_days: int = 365 * 2  # ~2 years
+    lookback_days: int = 730          # ~2 years
     sma_fast: int = 10
     sma_slow: int = 20
     rsi_period: int = 14
@@ -42,67 +44,148 @@ class FeatureSpec:
 
 def _rsi(close: pd.Series, period: int = 14) -> pd.Series:
     delta = close.diff()
-    gain = (delta.where(delta > 0, 0)).rolling(period).mean()
-    loss = (-delta.where(delta < 0, 0)).rolling(period).mean()
+    gain = (delta.where(delta > 0, 0.0)).rolling(period).mean()
+    loss = (-delta.where(delta < 0, 0.0)).rolling(period).mean()
     rs = gain / (loss + 1e-9)
     return 100 - (100 / (1 + rs))
 
 
-@lru_cache(maxsize=64)
-def fetch_prices_stooq(ticker: str) -> pd.DataFrame:
+def _normalize_stooq_ticker(ticker: str) -> str:
     """
-    Fetch daily OHLCV from Stooq (free). Returns ascending date index.
+    Stooq expects symbols like:
+      - aapl.us (US equities)
+      - msft.us
+    If user types AAPL, we convert to aapl.us.
+    If user already provides suffix (contains '.'), we keep it.
     """
-    # Stooq expects lowercase tickers for US like "aapl" (it also works with uppercase sometimes).
     t = ticker.strip().lower()
-    df = pdr.DataReader(t, "stooq")  # newest->oldest
-    df = df.sort_index()            # oldest->newest
-    df.index = pd.to_datetime(df.index)
+    if not t:
+        return t
+    if "." not in t:
+        # default: US equity
+        t = f"{t}.us"
+    return t
+
+
+@lru_cache(maxsize=128)
+def fetch_prices_stooq(ticker: str) -> pd.DataFrame:
+    """
+    Fetch daily OHLCV from Stooq via CSV.
+    Returns DataFrame indexed by Date ascending with columns:
+    Open, High, Low, Close, Volume
+    """
+    sym = _normalize_stooq_ticker(ticker)
+    url = f"https://stooq.com/q/d/l/?s={sym}&i=d"
+    r = requests.get(url, timeout=25)
+    r.raise_for_status()
+
+    df = pd.read_csv(io.StringIO(r.text))
+    if df.empty or "Date" not in df.columns:
+        raise ValueError(f"No data returned for ticker '{ticker}' (stooq symbol '{sym}').")
+
+    df["Date"] = pd.to_datetime(df["Date"])
+    df = df.set_index("Date").sort_index()
+
+    # Basic validation
+    needed = {"Open", "High", "Low", "Close"}
+    if not needed.issubset(set(df.columns)):
+        raise ValueError(f"Unexpected Stooq columns for '{ticker}': {list(df.columns)}")
+
+    # Ensure numeric
+    for c in ["Open", "High", "Low", "Close", "Volume"]:
+        if c in df.columns:
+            df[c] = pd.to_numeric(df[c], errors="coerce")
+
+    df = df.dropna(subset=["Close"]).copy()
     return df
 
 
-def build_features(df: pd.DataFrame, spec: FeatureSpec) -> pd.DataFrame:
+def build_features(prices: pd.DataFrame, spec: FeatureSpec) -> pd.DataFrame:
     """
-    Build simple features + binary target:
+    Build simple features + target:
       target = 1 if next-day return > 0 else 0
+    Keep ret_next for a toy backtest.
     """
-    out = df.copy()
-    out["close"] = out["Close"].astype(float)
+    df = prices.copy()
+    df["close"] = df["Close"].astype(float)
 
-    out["ret_1"] = out["close"].pct_change()
-    out["ret_5"] = out["close"].pct_change(5)
-    out["sma_fast"] = out["close"].rolling(spec.sma_fast).mean()
-    out["sma_slow"] = out["close"].rolling(spec.sma_slow).mean()
-    out["sma_ratio"] = out["sma_fast"] / (out["sma_slow"] + 1e-9) - 1.0
+    df["ret_1"] = df["close"].pct_change()
+    df["ret_5"] = df["close"].pct_change(5)
 
-    out["rsi"] = _rsi(out["close"], spec.rsi_period)
-    out["vol"] = out["ret_1"].rolling(spec.vol_window).std()
+    df["sma_fast"] = df["close"].rolling(spec.sma_fast).mean()
+    df["sma_slow"] = df["close"].rolling(spec.sma_slow).mean()
+    df["sma_ratio"] = df["sma_fast"] / (df["sma_slow"] + 1e-9) - 1.0
 
-    # next-day return and target
-    out["ret_next"] = out["close"].pct_change().shift(-1)
-    out["target"] = (out["ret_next"] > 0).astype(int)
+    df["rsi"] = _rsi(df["close"], spec.rsi_period)
+    df["vol"] = df["ret_1"].rolling(spec.vol_window).std()
 
-    # drop rows with NaNs from rolling calcs
-    out = out.dropna().copy()
+    # Next-day realized return and label
+    df["ret_next"] = df["close"].pct_change().shift(-1)
+    df["target"] = (df["ret_next"] > 0).astype(int)
 
-    # select model columns
-    feats = out[["ret_1", "ret_5", "sma_ratio", "rsi", "vol"]].copy()
-    feats["target"] = out["target"].astype(int).values
-    feats["close"] = out["close"].values
-    return feats
+    df = df.dropna().copy()
 
+    # Final dataset columns used by model + extras
+    out = df[["close", "ret_1", "ret_5", "sma_ratio", "rsi", "vol", "ret_next", "target"]].copy()
+    return out
 
-def make_dataset_for_tickers(tickers, spec: FeatureSpec) -> pd.DataFrame:
+
+def make_dataset_for_tickers(tickers: List[str], spec: FeatureSpec) -> Tuple[pd.DataFrame, List[str]]:
+    """
+    Returns:
+      - combined dataset with columns:
+          date, ticker, close, ret_1, ret_5, sma_ratio, rsi, vol, ret_next, target
+      - list of tickers that failed
+    """
     frames = []
+    failed = []
+
     for t in tickers:
-        prices = fetch_prices_stooq(t)
-        # keep last N days (plus some buffer for rolling windows)
-        prices = prices.iloc[-(spec.lookback_days + 100):].copy()
-        feats = build_features(prices, spec)
-        feats["ticker"] = t.upper()
-        feats["date"] = feats.index.astype(str)
-        frames.append(feats.reset_index(drop=True))
-    return pd.concat(frames, ignore_index=True)
+        try:
+            prices = fetch_prices_stooq(t)
+            # keep a window + buffer for rolling indicators
+            prices = prices.iloc[-(spec.lookback_days + 120):].copy()
+            feats = build_features(prices, spec)
+            feats = feats.reset_index().rename(columns={"Date": "date"})
+            feats["ticker"] = t.upper()
+            frames.append(feats)
+        except Exception:
+            failed.append(t.upper())
+
+    if not frames:
+        raise ValueError("No tickers returned usable data. Try different tickers (e.g., AAPL, MSFT).")
+
+    df = pd.concat(frames, ignore_index=True)
+    df["date"] = pd.to_datetime(df["date"])
+    df = df.sort_values(["ticker", "date"]).reset_index(drop=True)
+    return df, failed
+
+
+def split_train_val_per_ticker(df: pd.DataFrame, train_frac: float = 0.8) -> pd.DataFrame:
+    parts = []
+    for t, dft in df.groupby("ticker", sort=False):
+        dft = dft.sort_values("date").reset_index(drop=True)
+        n = len(dft)
+        cut = max(int(n * train_frac), 1)
+        dft["split"] = "train"
+        if cut < n:
+            dft.loc[cut:, "split"] = "val"
+        parts.append(dft)
+    return pd.concat(parts, ignore_index=True)
+
+
+def fig_to_image(fig) -> np.ndarray:
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", dpi=160)
+    plt.close(fig)
+    buf.seek(0)
+    return plt.imread(buf)
+
+
+def save_df_to_temp_csv(df: pd.DataFrame) -> str:
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="signals_dataset_")
+    df.to_csv(tmp.name, index=False)
+    return tmp.name
 
 
 # -----------------------------
@@ -140,25 +223,7 @@ class LitClassifier(pl.LightningModule):
         return torch.optim.Adam(self.parameters(), lr=self.lr)
 
 
-def fig_to_image(fig) -> np.ndarray:
-    buf = io.BytesIO()
-    fig.savefig(buf, format="png", bbox_inches="tight", dpi=160)
-    plt.close(fig)
-    buf.seek(0)
-    return plt.imread(buf)
-
-
-def save_df_to_temp_csv(df: pd.DataFrame) -> str:
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="signals_dataset_")
-    df.to_csv(tmp.name, index=False)
-    return tmp.name
-
-
-# -----------------------------
-# “Signals” logic (educational)
-# -----------------------------
 def signal_from_prob(p_up: float, buy_th: float, sell_th: float) -> str:
-    # p_up: probability next day is up
     if p_up >= buy_th:
         return "BUY (signal)"
     if p_up <= sell_th:
@@ -166,6 +231,9 @@ def signal_from_prob(p_up: float, buy_th: float, sell_th: float) -> str:
     return "HOLD (signal)"
 
 
+# -----------------------------
+# Main Gradio function
+# -----------------------------
 def run_app(
     tickers_text: str,
     lookback_days: int,
@@ -175,40 +243,31 @@ def run_app(
     seed: int,
     buy_threshold: float,
     sell_threshold: float,
+    device_choice: str,
 ):
     pl.seed_everything(int(seed), workers=True)
 
     tickers = [t.strip().upper() for t in tickers_text.split(",") if t.strip()]
     tickers = tickers[:10]
-    if len(tickers) == 0:
-        raise gr.Error("Please enter at least 1 ticker (comma-separated), e.g. AAPL, MSFT, NVDA")
+    if not tickers:
+        raise gr.Error("Enter at least 1 ticker, e.g. AAPL, MSFT, NVDA")
 
     spec = FeatureSpec(lookback_days=int(lookback_days))
+    df_raw, failed = make_dataset_for_tickers(tickers, spec)
 
-    # Build dataset (one table for all tickers)
-    df = make_dataset_for_tickers(tuple(tickers), spec)  # tuple so cache plays nice if you reuse
-    df = df.replace([np.inf, -np.inf], np.nan).dropna().copy()
-
-    # Train/val split by time-ish: last 20% as val PER ticker
-    parts = []
-    for t in tickers:
-        dft = df[df["ticker"] == t].copy()
-        n = len(dft)
-        cut = max(int(n * 0.8), 1)
-        dft["split"] = "train"
-        dft.loc[dft.index[cut:], "split"] = "val"
-        parts.append(dft)
-    df = pd.concat(parts, ignore_index=True)
+    # split per ticker
+    df = split_train_val_per_ticker(df_raw, train_frac=0.8)
 
     feature_cols = ["ret_1", "ret_5", "sma_ratio", "rsi", "vol"]
 
-    # Standardize features using TRAIN split stats
+    # Standardize using TRAIN split stats
     train_df = df[df["split"] == "train"].copy()
     mu = train_df[feature_cols].mean()
     sd = train_df[feature_cols].std().replace(0, 1.0)
 
     df_std = df.copy()
     df_std[feature_cols] = (df_std[feature_cols] - mu) / sd
+    df_std = df_std.replace([np.inf, -np.inf], np.nan).dropna().copy()
 
     # Torch tensors
     X_train = torch.tensor(df_std[df_std["split"] == "train"][feature_cols].values, dtype=torch.float32)
@@ -220,10 +279,18 @@ def run_app(
     train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=int(batch_size), shuffle=True)
     val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=int(batch_size), shuffle=False)
 
+    # Lightning Trainer device selection
+    want_cuda = (device_choice == "cuda")
+    has_cuda = torch.cuda.is_available()
+    using_cuda = want_cuda and has_cuda
+    accelerator = "gpu" if using_cuda else "cpu"
+
     model = LitClassifier(n_features=len(feature_cols), lr=float(lr))
 
     trainer = pl.Trainer(
         max_epochs=int(epochs),
+        accelerator=accelerator,
+        devices=1,
         logger=False,
         enable_checkpointing=False,
         enable_progress_bar=False,
@@ -233,11 +300,13 @@ def run_app(
     trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)
 
     # Inference: latest row per ticker
-    out_rows = []
     model.eval()
+    out_rows = []
     with torch.no_grad():
         for t in tickers:
-            dft = df_std[df_std["ticker"] == t].copy()
+            dft = df_std[df_std["ticker"] == t].sort_values("date")
+            if dft.empty:
+                continue
             last = dft.iloc[-1]
             x = torch.tensor(last[feature_cols].values.astype(np.float32)).unsqueeze(0)
             logit = model(x).item()
@@ -246,63 +315,73 @@ def run_app(
             out_rows.append(
                 {
                     "ticker": t,
-                    "last_close": float(last["close"]),
+                    "date": last["date"].date().isoformat(),
+                    "last_close": round(float(last["close"]), 4),
                     "p_up_next_day": round(float(p_up), 4),
                     "signal": sig,
-                    "date": last["date"],
                 }
             )
-    signals_df = pd.DataFrame(out_rows).sort_values("p_up_next_day", ascending=False)
 
-    # Simple backtest curve for FIRST ticker (val split only): long if p_up>=buy_th, short if p_up<=sell_th
+    signals_df = pd.DataFrame(out_rows)
+    if not signals_df.empty:
+        signals_df = signals_df.sort_values("p_up_next_day", ascending=False).reset_index(drop=True)
+
+    # Toy backtest for first ticker (val split only)
+    backtest_img = None
     t0 = tickers[0]
-    d0 = df_std[df_std["ticker"] == t0].copy()
-    d0_val = d0[d0["split"] == "val"].copy()
-    if len(d0_val) >= 5:
-        X0 = torch.tensor(d0_val[feature_cols].values, dtype=torch.float32)
-        logits = model(X0).detach().cpu().numpy()
+    d0 = df_std[(df_std["ticker"] == t0) & (df_std["split"] == "val")].sort_values("date").copy()
+    if len(d0) >= 30:
+        X0 = torch.tensor(d0[feature_cols].values, dtype=torch.float32)
+        with torch.no_grad():
+            logits = model(X0).detach().cpu().numpy()
         p = 1 / (1 + np.exp(-logits))
 
-        pos = np.zeros_like(p)
+        pos = np.zeros_like(p, dtype=float)
         pos[p >= float(buy_threshold)] = 1.0
         pos[p <= float(sell_threshold)] = -1.0  # short (toy)
-        # realized return is ret_next (already aligned)
-        r = d0_val["ret_next"].values
-        strat = pos * r
+
+        strat = pos * d0["ret_next"].values
         equity = (1 + strat).cumprod()
 
         fig = plt.figure()
         plt.plot(equity)
-        plt.title(f"Toy Backtest (VAL only) — {t0}  | long/short by signal")
+        plt.title(f"Toy Backtest (VAL only) — {t0} | long/short by signal")
         plt.xlabel("Val days")
         plt.ylabel("Equity (start=1.0)")
         plt.grid(True, alpha=0.3)
         backtest_img = fig_to_image(fig)
-    else:
-        backtest_img = None
-
-    # Data preview + download
-    preview = df.head(20).copy()
-    csv_path = save_df_to_temp_csv(df)
-
-    summary = (
-        f"Tickers: {', '.join(tickers)}\n"
-        f"Rows total: {len(df)} (train={len(df[df['split']=='train'])}, val={len(df[df['split']=='val'])})\n"
-        f"Model: Lightning MLP classifier (predict next-day up/down)\n"
-        f"Signals are educational thresholds: BUY if p>= {buy_threshold}, SELL if p<= {sell_threshold}\n"
-    )
 
-    return signals_df, backtest_img, preview, csv_path, summary
+    # Data preview + download (download the unstandardized feature table with split)
+    export_df = df.copy()
+    export_df["date"] = export_df["date"].dt.date.astype(str)
+    export_df = export_df[["date", "ticker", "split", "close", "ret_1", "ret_5", "sma_ratio", "rsi", "vol", "ret_next", "target"]]
+    preview_df = export_df.head(25).round(6)
+    csv_path = save_df_to_temp_csv(export_df.round(8))
+
+    summary_lines = [
+        f"Using device: {'cuda' if using_cuda else 'cpu'}",
+        f"Tickers requested (max 10): {', '.join(tickers)}",
+        f"Rows: {len(export_df)} | train={int((export_df['split']=='train').sum())} | val={int((export_df['split']=='val').sum())}",
+        f"BUY if p_up >= {buy_threshold:.2f} | SELL if p_up <= {sell_threshold:.2f}",
+    ]
+    if failed:
+        summary_lines.append(f"Tickers with no data / error: {', '.join(failed)}")
+    summary = "\n".join(summary_lines)
 
+    return signals_df, backtest_img, preview_df, csv_path, summary
 
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
 with gr.Blocks(title="Educational Stock Signals (Lightning)") as demo:
-    gr.Markdown("# Educational Stock Signals (Top 10 tickers)\n" + DISCLAIMER)
+    gr.Markdown("# Educational Stock Signals (Lightning)\n" + DISCLAIMER)
 
-    with gr.Row():
-        tickers_text = gr.Textbox(
-            value="AAPL, MSFT, NVDA, AMZN, GOOGL, META, TSLA, JPM, V, XOM",
-            label="Tickers (comma-separated, up to 10)",
-        )
+    tickers_text = gr.Textbox(
+        value="AAPL, MSFT, NVDA, AMZN, GOOGL, META, TSLA, JPM, V, XOM",
+        label="Tickers (comma-separated, up to 10)",
+        info="Tip: Stooq uses US symbols like AAPL -> aapl.us automatically. If needed, specify suffix (e.g., '7203.jp').",
+    )
 
     with gr.Row():
         lookback_days = gr.Slider(200, 2000, value=730, step=10, label="Lookback days (history window)")
@@ -316,24 +395,28 @@ with gr.Blocks(title="Educational Stock Signals (Lightning)") as demo:
     with gr.Row():
         buy_threshold = gr.Slider(0.50, 0.80, value=0.55, step=0.01, label="BUY threshold (p_up)")
         sell_threshold = gr.Slider(0.20, 0.50, value=0.45, step=0.01, label="SELL threshold (p_up)")
+        device_choice = gr.Radio(["cpu", "cuda"], value="cpu", label="Device (cuda only if available)")
 
     run_btn = gr.Button("Build signals", variant="primary")
 
     with gr.Tab("Signals"):
-        signals_df = gr.Dataframe(label="Signals (educational)", wrap=True)
+        signals_out = gr.Dataframe(label="Signals (educational)", wrap=True)
 
     with gr.Tab("Backtest (toy)"):
-        backtest_img = gr.Image(label="Toy equity curve (val only; first ticker)", type="numpy")
+        backtest_out = gr.Image(label="Toy equity curve (val only; first ticker)", type="numpy")
 
     with gr.Tab("Data"):
-        preview_df = gr.Dataframe(label="Feature dataset preview (first 20 rows)", wrap=True)
-        download_file = gr.File(label="Download full dataset CSV (features + target)")
-        summary_txt = gr.Textbox(label="Run summary", lines=6)
+        preview_out = gr.Dataframe(label="Feature dataset preview", wrap=True)
+        download_out = gr.File(label="Download full dataset CSV (features + target + split)")
+        summary_out = gr.Textbox(label="Run summary", lines=8)
 
     run_btn.click(
         fn=run_app,
-        inputs=[tickers_text, lookback_days, lr, batch_size, epochs, seed, buy_threshold, sell_threshold],
-        outputs=[signals_df, backtest_img, preview_df, download_file, summary_txt],
+        inputs=[
+            tickers_text, lookback_days, lr, batch_size, epochs, seed,
+            buy_threshold, sell_threshold, device_choice
+        ],
+        outputs=[signals_out, backtest_out, preview_out, download_out, summary_out],
     )
 
 if __name__ == "__main__":