File size: 4,558 Bytes
b54319d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | import os
import urllib.request
import pandas as pd
import numpy as np
from typing import List, Optional, Tuple
class HuggingFaceETFLoader:
"""
Data Loader for ETF price data (Hugging Face / yfinance fallback).
Downloads and caches parquet files locally in `data/raw/` for offline reproducibility.
"""
DEFAULT_UNIVERSE = [
"SPY", "QQQ", "IWM", "MDY",
"XLK", "XLF", "XLE", "XLV", "XLY", "XLP", "XLI", "XLB", "XLU", "XLC", "XLRE",
"SMH", "XBI", "KRE", "ITB",
"TLT", "IEF", "SHY", "LQD", "HYG", "TIP",
"GLD", "SLV", "USO", "DBA",
"EEM", "EFA", "FXI", "EWJ",
"MTUM", "USMV", "QUAL", "IWD", "IWF"
]
def __init__(self, cache_dir: str = "data/raw"):
self.cache_dir = cache_dir
os.makedirs(self.cache_dir, exist_ok=True)
def load_prices(self, tickers: Optional[List[str]] = None) -> pd.DataFrame:
"""Load ETF prices DataFrame. Uses cached parquet if available, otherwise downloads via yfinance."""
local_path = os.path.join(self.cache_dir, "hf_prices.parquet")
if os.path.exists(local_path):
print(f"[ETFLoader] Reading cached dataset from {local_path}...")
df = pd.read_parquet(local_path)
df["date"] = pd.to_datetime(df["date"])
return df
return self._fetch_yfinance_prices(tickers)
def _fetch_yfinance_prices(self, tickers: Optional[List[str]] = None) -> pd.DataFrame:
"""Fetch daily OHLCV dataset using yfinance and save to local parquet cache."""
import yfinance as yf
tickers = tickers or self.DEFAULT_UNIVERSE
print(f"[ETFLoader] Fetching {len(tickers)} tickers via yfinance...")
try:
data = yf.download(tickers, start="2014-01-01", end="2024-01-01", group_by="ticker", auto_adjust=False, progress=False)
records = []
for t in tickers:
if t in data.columns.levels[0]:
sub = data[t].dropna(how="all").reset_index()
if len(sub) > 0:
sub["symbol"] = t
sub.rename(columns={
"Date": "date", "Open": "open", "High": "high", "Low": "low",
"Close": "close", "Adj Close": "adjusted_close", "Volume": "volume"
}, inplace=True)
records.append(sub)
if records:
df_res = pd.concat(records, ignore_index=True)
df_res["date"] = pd.to_datetime(df_res["date"])
if "adjusted_close" not in df_res.columns:
df_res["adjusted_close"] = df_res["close"]
# Save cache
df_res.to_parquet(os.path.join(self.cache_dir, "hf_prices.parquet"))
print(f"[ETFLoader] Cached {len(df_res)} price rows to {self.cache_dir}/hf_prices.parquet")
return df_res
except Exception as e:
print(f"[ETFLoader] Remote download failed: {e}")
# Fallback to synthetic offline dataset generator
return self.generate_synthetic_prices(tickers)
@staticmethod
def generate_synthetic_prices(tickers: List[str], num_days: int = 1000) -> pd.DataFrame:
"""Generate realistic synthetic OHLCV data for offline CI / local smoke tests."""
print("[ETFLoader] Generating synthetic price data for offline testing...")
dates = pd.date_range("2020-01-01", periods=num_days, freq="B")
records = []
np.random.seed(42)
for t in tickers:
price = 100.0 + np.random.uniform(-10, 10)
volatility = 0.015 + np.random.uniform(0.005, 0.01)
prices = [price]
for _ in range(num_days - 1):
ret = np.random.normal(0.0004, volatility)
price *= np.exp(ret)
prices.append(price)
df_t = pd.DataFrame({
"date": dates,
"symbol": t,
"open": prices,
"high": [p * (1 + abs(np.random.normal(0, 0.005))) for p in prices],
"low": [p * (1 - abs(np.random.normal(0, 0.005))) for p in prices],
"close": prices,
"adjusted_close": prices,
"volume": np.random.uniform(1e6, 1e7, size=num_days)
})
records.append(df_t)
return pd.concat(records, ignore_index=True)
|