Spaces:
Sleeping
Sleeping
| import datetime | |
| import yfinance as yf | |
| import pandas as pd | |
| import logging | |
| from yfinance.search import Search | |
| import numpy as np | |
| class SP500Sector: | |
| MATERIALS = "Materials" | |
| COMMUNICATION_SERVICES = "Communication Services" | |
| CONSUMER_DISCRETIONARY = "Consumer Discretionary" | |
| CONSUMER_STAPLES = "Consumer Staples" | |
| ENERGY = "Energy" | |
| FINANCIALS = "Financials" | |
| HEALTH_CARE = "Health Care" | |
| INDUSTRIALS = "Industrials" | |
| OTHER = "Taiwan Other" | |
| REAL_ESTATE = "Real Estate" | |
| INFORMATION_TECHNOLOGY = "Information Technology" | |
| UTILITIES = "Utilities" | |
| def get_all_sectors(cls): | |
| return sorted([v for k, v in vars(cls).items() if k.isupper() and isinstance(v, str)]) | |
| def get_sector_id(cls, sector_str): | |
| sectors = cls.get_all_sectors() | |
| if sector_str in sectors: | |
| return sectors.index(sector_str) | |
| return sectors.index(cls.OTHER) | |
| def isMember(self, sector_str): | |
| return sector_str in self.get_all_sectors() | |
| def fromYfinanceSector(self, sector_str): | |
| if sector_str == "Basic Materials": | |
| return self.MATERIALS | |
| elif sector_str == "Communication Services": | |
| return self.COMMUNICATION_SERVICES | |
| elif sector_str == "Consumer Cyclical": | |
| return self.CONSUMER_DISCRETIONARY | |
| elif sector_str == "Consumer Defensive": | |
| return self.CONSUMER_STAPLES | |
| elif sector_str == "Energy": | |
| return self.ENERGY | |
| elif sector_str == "Financial Services": | |
| return self.FINANCIALS | |
| elif sector_str == "Healthcare": | |
| return self.HEALTH_CARE | |
| elif sector_str == "Industrials": | |
| return self.INDUSTRIALS | |
| elif sector_str == "Real Estate": | |
| return self.REAL_ESTATE | |
| elif sector_str == "Technology": | |
| return self.INFORMATION_TECHNOLOGY | |
| elif sector_str == "Utilities": | |
| return self.UTILITIES | |
| else: | |
| return self.OTHER | |
| def get_yfinance_ticker_info(code) -> dict: | |
| """ | |
| Get ticker information using yfinance search method, which is faster than downloading history. | |
| Returns a dictionary with keys | |
| """ | |
| if not hasattr(get_yfinance_ticker_info, "call_counter"): | |
| get_yfinance_ticker_info.call_counter = 0 | |
| get_yfinance_ticker_info.call_counter += 1 | |
| # Suppress yfinance logging to avoid cluttering output | |
| logger = logging.getLogger('yfinance') | |
| original_level = logger.level | |
| logger.setLevel(logging.CRITICAL) | |
| try: | |
| call_count = get_yfinance_ticker_info.call_counter | |
| results = Search(code, max_results=10).quotes | |
| if call_count == 1: | |
| print(results[0]) | |
| print(f"{get_yfinance_ticker_info.call_counter}. Searched results for '{code}': {[x.get('symbol') for x in results]}") | |
| for r in results: | |
| if r.get("symbol") == code: | |
| return {k: r.get(k, '') for k in ["exchange", "symbol", "longname", "sector"]} | |
| for r in results: | |
| if r.get("symbol", "").startswith(code): | |
| return {k: r.get(k, '') for k in ["exchange", "symbol", "longname", "sector"]} | |
| if results: | |
| r = results[0] | |
| return {k: r.get(k, '') for k in ["exchange", "symbol", "longname", "sector"]} | |
| return None | |
| except Exception as e: | |
| print(f"Warning: Search method failed for '{code}' with error: {e}.") | |
| finally: | |
| logger.setLevel(original_level) | |
| print(f"Warning: Could not find valid ticker for {code} using search.") | |
| return None | |
| def get_ticker_start_date(code): | |
| """ | |
| Get the start date (first trade date) of a stock without downloading all data. | |
| """ | |
| try: | |
| ticker = yf.Ticker(code) | |
| # get_history_metadata fetches lightweight metadata about the listing | |
| meta = ticker.get_history_metadata() | |
| return meta.get("firstTradeDate", None) | |
| except Exception as e: | |
| print(f"Warning: Could not fetch start date for '{code}': {e}") | |
| return None | |
| def calculate_return(df: pd.DataFrame) -> pd.Series: | |
| return df['Close'].pct_change() | |
| def calculate_log_return(df: pd.DataFrame) -> pd.Series: | |
| return np.log(df['Close'] / df['Close'].shift(1)) | |
| def calculate_ma_5(df: pd.DataFrame) -> pd.Series: | |
| return df['Close'].rolling(window=5).mean() | |
| def calculate_ma_20(df: pd.DataFrame) -> pd.Series: | |
| return df['Close'].rolling(window=20).mean() | |
| def calculate_vol_20(df: pd.DataFrame) -> pd.Series: | |
| return df['Return'].rolling(window=20).std() | |
| def calculate_roc_10(df: pd.DataFrame) -> pd.Series: | |
| return (df['Close'] - df['Close'].shift(10)) / df['Close'].shift(10) | |
| def calculate_rsi_14(df: pd.DataFrame) -> pd.Series: | |
| window = 14 | |
| delta = df['Close'].diff() | |
| gain = (delta.where(delta > 0, 0)).rolling(window=window).mean() | |
| loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean() | |
| rs = gain / loss | |
| return 100 - (100 / (1 + rs)) | |
| def calculate_macd(df: pd.DataFrame) -> pd.Series: | |
| ema_12 = df['Close'].ewm(span=12, adjust=False).mean() | |
| ema_26 = df['Close'].ewm(span=26, adjust=False).mean() | |
| return ema_12 - ema_26 | |
| def calculate_macd_signal(df: pd.DataFrame) -> pd.Series: | |
| macd = calculate_macd(df) | |
| return macd.ewm(span=9, adjust=False).mean() | |
| def calculate_kd_k(df: pd.DataFrame, n: int = 9) -> pd.Series: | |
| if 'High' not in df.columns or 'Low' not in df.columns: | |
| return pd.Series(50.0, index=df.index) | |
| low_min = df['Low'].rolling(window=n).min() | |
| high_max = df['High'].rolling(window=n).max() | |
| rsv = 100 * (df['Close'] - low_min) / (high_max - low_min + 1e-8) | |
| # com=2 is exponentially equivalent to alpha=1/3, matching standard KD smoothing | |
| return rsv.ewm(com=2, adjust=False).mean() | |
| def calculate_kd_d(df: pd.DataFrame) -> pd.Series: | |
| k = calculate_kd_k(df) | |
| return k.ewm(com=2, adjust=False).mean() | |
| def calculate_bb_upper(df: pd.DataFrame) -> pd.Series: | |
| ma_20 = calculate_ma_20(df) | |
| std_20 = df['Close'].rolling(window=20).std() | |
| return ma_20 + 2 * std_20 | |
| def calculate_bb_lower(df: pd.DataFrame) -> pd.Series: | |
| ma_20 = calculate_ma_20(df) | |
| std_20 = df['Close'].rolling(window=20).std() | |
| return ma_20 - 2 * std_20 | |
| def calculate_bias_6(df: pd.DataFrame) -> pd.Series: | |
| ma_6 = df['Close'].rolling(window=6).mean() | |
| return (df['Close'] - ma_6) / ma_6 * 100 | |
| feature_map = { | |
| 'Return': calculate_return, | |
| 'Log_Return': calculate_log_return, | |
| 'MA_5': calculate_ma_5, | |
| 'MA_20': calculate_ma_20, | |
| 'Vol_20': calculate_vol_20, | |
| 'ROC_10': calculate_roc_10, | |
| 'RSI_14': calculate_rsi_14, | |
| 'MACD': calculate_macd, | |
| 'MACD_Signal': calculate_macd_signal, | |
| 'KD_K': calculate_kd_k, | |
| 'KD_D': calculate_kd_d, | |
| 'BB_Upper': calculate_bb_upper, | |
| 'BB_Lower': calculate_bb_lower, | |
| 'BIAS_6': calculate_bias_6 | |
| } | |
| def generate_features(df: pd.DataFrame, sector_str: str = SP500Sector.OTHER, is_training: bool = True) -> pd.DataFrame: | |
| for feature_name, calc_func in feature_map.items(): | |
| df[feature_name] = calc_func(df) | |
| df['Sector_ID'] = SP500Sector.get_sector_id(sector_str) | |
| if is_training: | |
| df['Target'] = df['Close'].shift(-1) / df['Close'] - 1.0 | |
| return df | |