feat(src): initialize a framework (#6)

lwaekfjlk · web-flow · commit 151828283f9a · 2025-06-25T17:46:47.000-07:00
* update

* update
diff --git a/data/yfinance_data/price_data/AAPL_data_formatted.json b/data/yfinance_data/price_data/AAPL_data_formatted.json
diff --git a/examples/demo.py b/examples/demo.py
@@ -0,0 +1,52 @@
+import argparse
+
+from trading_bench.bench import SimBench
+from trading_bench.model_wrapper import RuleBasedModel
+from trading_bench.utils import setup_logging
+
+
+def main():
+    setup_logging()
+    parser = argparse.ArgumentParser(
+        description='Run the simulated trading bench with Finnhub data crawler'
+    )
+    parser.add_argument('--ticker', required=True, help='Stock ticker symbol')
+    parser.add_argument('--start_date', required=True, help='Start date (YYYY-MM-DD)')
+    parser.add_argument('--end_date', required=True, help='End date (YYYY-MM-DD)')
+    parser.add_argument(
+        '--data_dir', required=True, help='Root directory where data is stored'
+    )
+    parser.add_argument(
+        '--eval_delay', type=int, default=5, help='Evaluation delay in data points'
+    )
+    parser.add_argument(
+        '--resolution',
+        choices=['1', '5', '15', '30', '60', 'D', 'W', 'M'],
+        default='D',
+        help="Finnhub data resolution (e.g., '1','5','15','30','60','D','W','M')",
+    )
+
+    args = parser.parse_args()
+
+    model = RuleBasedModel()
+    bench = SimBench(
+        ticker=args.ticker,
+        start_date=args.start_date,
+        end_date=args.end_date,
+        data_dir=args.data_dir,
+        model=model,
+        eval_delay=args.eval_delay,
+        resolution=args.resolution,
+    )
+    summary = bench.run()
+
+    print('Performance Summary:')
+    for key, value in summary.items():
+        if isinstance(value, float):
+            print(f'{key}: {value:.4f}')
+        else:
+            print(f'{key}: {value}')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,19 +1,18 @@
 [tool.poetry]
-name = "python-project-template"
+name = "trading-bench"
 version = "0.0.1"
 description = "A template for python-based research project"
 authors = ["Haofei Yu <yuhaofei44@gmail.com>"]
 license = "Apache 2.0 License"
 readme = "README.md"
-packages = [
-    { include = "src" }
-]
 
 [tool.poetry.dependencies]
 python = ">=3.9, <3.12"
 mypy = "^1.8.0"
 beartype = "^0.17.1"
 pydantic = "^2.8.2"
+finnhub-python = "*"
+yfinance = "*"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.6.0"
diff --git a/scripts/demo.sh b/scripts/demo.sh
@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+
+# === Configuration ===
+DATA_DIR="../data"    # e.g. ~/projects/data
+TICKER="AAPL"
+START_DATE="2025-01-01"
+END_DATE="2025-06-01"
+EVAL_DELAY=5
+RESOLUTION="D"        # Finnhub resolution: 1,5,15,30,60,D,W,M
+
+# === Run the bench ===
+python ../examples/demo.py  \
+  --ticker      "${TICKER}" \
+  --start_date  "${START_DATE}" \
+  --end_date    "${END_DATE}" \
+  --data_dir    "${DATA_DIR}" \
+  --eval_delay  "${EVAL_DELAY}" \
+  --resolution  "${RESOLUTION}"
diff --git a/trading_bench/.gitkeep b/trading_bench/.gitkeep
diff --git a/trading_bench/__init__.py b/trading_bench/__init__.py
diff --git a/trading_bench/bench.py b/trading_bench/bench.py
@@ -0,0 +1,103 @@
+import json
+import os
+from collections import defaultdict, deque
+from datetime import datetime
+
+from .data_fetcher import fetch_price_data
+from .evaluator import ReturnEvaluator
+from .metrics import MetricsLogger
+from .model_wrapper import BaseModel
+from .signal import Signal
+
+
+class SimBench:
+    """
+    Simulated backtest bench that asks a model for buy signals and evaluates returns.
+    Uses the updated yfinance-based fetch_price_data to retrieve OHLCV data.
+    """
+
+    def __init__(
+        self,
+        ticker: str,
+        start_date: str,
+        end_date: str,
+        data_dir: str,
+        model: BaseModel,
+        eval_delay: int = 5,
+        resolution: str = 'D',
+    ):
+        self.ticker = ticker
+        self.start_date = start_date
+        self.end_date = end_date
+        self.data_dir = data_dir
+        self.model = model
+        self.eval_delay = eval_delay
+        self.resolution = resolution
+
+        # Fetch and save price data (yfinance) into yfinance_data/price_data
+        fetch_price_data(
+            ticker=self.ticker,
+            start_date=self.start_date,
+            end_date=self.end_date,
+            data_dir=self.data_dir,
+            resolution=self.resolution,
+        )
+
+        # Load fetched JSON data from yfinance_data
+        data_path = os.path.join(
+            self.data_dir,
+            'yfinance_data',
+            'price_data',
+            f'{self.ticker}_data_formatted.json',
+        )
+        if not os.path.isfile(data_path):
+            raise FileNotFoundError(f'Expected data file not found at {data_path}')
+
+        with open(data_path, encoding='utf-8') as f:
+            raw_data = json.load(f)
+
+        # Parse into list of (datetime, close_price)
+        parsed: list[tuple[datetime, float]] = []
+        for date_str, v in raw_data.items():
+            date = datetime.fromisoformat(date_str)
+            # v is a dict with keys open, high, low, close, volume
+            price = float(v.get('close', v))
+            parsed.append((date, price))
+
+        # Sort chronologically and initialize history deque
+        self.data: list[tuple[datetime, float]] = sorted(parsed, key=lambda x: x[0])
+        self.history: deque = deque(self.data)
+
+        self.evaluator = ReturnEvaluator()
+        self.logger = MetricsLogger()
+
+    def run(self) -> dict[str, float]:
+        prices = [price for _, price in self.data]
+        n = len(prices)
+
+        # 1. start with an empty "past history" and a place to stash pending signals
+        self.history = deque()
+        pending: dict[int, list[Signal]] = defaultdict(list)
+
+        for idx, (date, price) in enumerate(self.data):
+            # 2. append this step into your history
+            self.history.append((date, price))
+
+            # 3. if model says BUY, schedule evaluation at idx + eval_delay
+            if self.model.should_buy([p for _, p in self.history]):
+                eval_idx = min(idx + self.eval_delay, n - 1)
+                eval_time = self.data[eval_idx][0]
+                signal = Signal(date, price, eval_time)
+                pending[eval_idx].append(signal)
+
+            # 4. now check if any scheduled signals are due at this idx
+            if idx in pending:
+                for signal in pending.pop(idx):
+                    # you might want to pass only the slice of history from buy→eval
+                    # but ReturnEvaluator could also just use signal.price + actual price at eval_time
+                    ret = self.evaluator.evaluate(
+                        signal, list(self.history)[-self.eval_delay - 1 :]
+                    )
+                    self.logger.record(ret)
+
+        return self.logger.summary()
diff --git a/trading_bench/config.py b/trading_bench/config.py
diff --git a/trading_bench/data_fetcher.py b/trading_bench/data_fetcher.py
@@ -0,0 +1,69 @@
+import json
+import os
+import time
+
+import yfinance as yf
+
+
+def fetch_price_data(
+    ticker: str, start_date: str, end_date: str, data_dir: str, resolution: str = 'D'
+) -> None:
+    """
+    Fetches historical OHLCV price data for a ticker via yfinance and saves it as formatted JSON.
+    Args:
+        ticker:     Stock ticker symbol.
+        start_date: YYYY-MM-DD
+        end_date:   YYYY-MM-DD
+        data_dir:   Root data directory where 'yfinance_data/price_data' will live.
+        resolution: '1', '5', '15', '30', '60', 'D', 'W', 'M'
+    """
+    # map your resolution codes to yfinance intervals
+    interval_map = {
+        '1': '1m',
+        '5': '5m',
+        '15': '15m',
+        '30': '30m',
+        '60': '60m',
+        'D': '1d',
+        'W': '1wk',
+        'M': '1mo',
+    }
+    interval = interval_map.get(resolution.upper(), '1d')
+
+    # download data
+    df = yf.download(
+        tickers=ticker,
+        start=start_date,
+        end=end_date,
+        interval=interval,
+        progress=False,
+    )
+
+    if df.empty:
+        raise RuntimeError(
+            f'No data returned for {ticker} {start_date}→{end_date} @ {interval}'
+        )
+
+    # Build date-indexed dict
+    data = {}
+    for idx, row in df.iterrows():
+        # idx is a pandas.Timestamp
+        date_str = idx.strftime('%Y-%m-%d')
+        data[date_str] = {
+            'open': float(row['Open']),
+            'high': float(row['High']),
+            'low': float(row['Low']),
+            'close': float(row['Close']),
+            'volume': int(row['Volume']),
+        }
+
+    # Ensure target directory exists
+    out_dir = os.path.join(data_dir, 'yfinance_data', 'price_data')
+    os.makedirs(out_dir, exist_ok=True)
+
+    out_path = os.path.join(out_dir, f'{ticker}_data_formatted.json')
+    with open(out_path, 'w', encoding='utf-8') as f:
+        json.dump(data, f, indent=2)
+
+    # be polite with any rate limits
+    time.sleep(1)
diff --git a/trading_bench/evaluator.py b/trading_bench/evaluator.py
@@ -0,0 +1,29 @@
+from collections import deque
+from datetime import datetime
+
+from .signal import Signal
+
+
+class ReturnEvaluator:
+    """
+    Evaluates return percentage for a given Signal based on recent price history.
+    """
+
+    def evaluate(self, signal: Signal, history: deque[tuple[datetime, float]]) -> float:
+        """
+        Find the first price in history at or after signal.eval_time and compute return.
+        If no such timestamp, use the latest price in history.
+
+        Returns:
+            return_pct: (exit_price - entry_price) / entry_price
+        """
+        exit_price = None
+        for ts, price in history:
+            if ts >= signal.eval_time:
+                exit_price = price
+                break
+        if exit_price is None:
+            # fallback to last price available
+            exit_price = history[-1][1]
+
+        return (exit_price - signal.entry_price) / signal.entry_price
diff --git a/trading_bench/metrics.py b/trading_bench/metrics.py
@@ -0,0 +1,43 @@
+class MetricsLogger:
+    """
+    Aggregates return percentages and provides summary statistics.
+    """
+
+    def __init__(self):
+        self.returns: list[float] = []
+
+    def record(self, return_pct: float) -> None:
+        """Add a realized return percentage from one trade."""
+        self.returns.append(return_pct)
+
+    def summary(self) -> dict[str, float]:
+        """
+        Compute summary metrics:
+          - trades: total number of trades
+          - avg_return: mean return
+          - win_rate: proportion of positive returns
+          - max_return: highest return
+          - min_return: lowest return
+        """
+        if not self.returns:
+            return {
+                'trades': 0,
+                'avg_return': 0.0,
+                'win_rate': 0.0,
+                'max_return': 0.0,
+                'min_return': 0.0,
+            }
+
+        trades = len(self.returns)
+        avg_return = sum(self.returns) / trades
+        win_rate = sum(1 for r in self.returns if r > 0) / trades
+        max_return = max(self.returns)
+        min_return = min(self.returns)
+
+        return {
+            'trades': trades,
+            'avg_return': avg_return,
+            'win_rate': win_rate,
+            'max_return': max_return,
+            'min_return': min_return,
+        }
diff --git a/trading_bench/model_wrapper.py b/trading_bench/model_wrapper.py
@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+
+
+class BaseModel(ABC):
+    @abstractmethod
+    def should_buy(self, price_history: list[float]) -> bool:
+        """Given recent price history, return True if the model signals a buy action."""
+        pass
+
+
+class RuleBasedModel(BaseModel):
+    """A simple rule-based model that buys when the latest price is below the historical average."""
+
+    def should_buy(self, price_history: list[float]) -> bool:
+        if not price_history:
+            return False
+        avg_price = sum(price_history) / len(price_history)
+        return price_history[-1] < avg_price
diff --git a/trading_bench/signal.py b/trading_bench/signal.py
@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from datetime import datetime
+
+
+@dataclass
+class Signal:
+    entry_time: datetime
+    entry_price: float
+    eval_time: datetime
+
+
+@dataclass
+class ReturnRecord:
+    entry_time: datetime
+    exit_time: datetime
+    return_pct: float
diff --git a/trading_bench/utils.py b/trading_bench/utils.py
@@ -0,0 +1,16 @@
+import logging
+from datetime import datetime
+
+
+def setup_logging(level=logging.INFO):
+    """
+    Configure root logger formatting and level.
+    """
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=level)
+
+
+def parse_date(date_str: str) -> datetime:
+    """
+    Parse an ISO-8601 date string into a datetime object.
+    """
+    return datetime.fromisoformat(date_str)
