quantmine 0.2.0__py3-none-any.whl

quantmine/data_acquisition.py

@@ -0,0 +1,326 @@
+import yfinance as yf
+import pandas as pd
+import time
+import json
+import glob
+import pandas as pd
+import os
+import hashlib
+import pandas_datareader.data as web
+from typing import Literal
+
+def is_delisted_error(msg: str, ignore_json: json)-> bool:
+    """Check whether an error message matches any delisting keyword.
+
+    Args:
+        msg: Error message to inspect.
+        ignore_json: Iterable of keyword strings used to detect delisting.
+
+    Returns:
+        True if any keyword appears in the message, otherwise False.
+    """
+    msg = msg.lower()
+    return any(kw in msg for kw in ignore_json)
+
+def load_blacklist(checkpoint_dir:str)-> set :
+    """Load the ticker blacklist from the checkpoint directory.
+
+    Args:
+        checkpoint_dir: Directory containing ``blacklist.json``.
+
+    Returns:
+        A set of blacklisted tickers. Returns an empty set if the file does not
+        exist.
+    """
+    path = os.path.join(checkpoint_dir, "blacklist.json")
+    if not os.path.exists(path):
+        return set()
+    with open(path) as f:
+        return set(json.load(f))
+    
+def get_ff3(start_date: str, end_date: str, save_path: str, format: Literal['csv', 'parquet']):
+    if format not in ('csv', 'parquet'):
+        raise ValueError(f'{format} format error')
+    start_date = pd.to_datetime(start_date)
+    end_date = pd.to_datetime(end_date)
+    output_path = os.path.join(os.getcwd(), save_path)
+    
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+
+    ff3 = web.DataReader('F-F_Research_Data_Factors', 'famafrench', start_date, end_date)
+    if format == 'csv':
+        ff3.to_csv(os.path.join(output_path, 'F-F_Research_Data_Factors.csv'))
+    if format == 'parquet':
+        ff3.to_parquet(os.path.join(output_path, 'F-F_Research_Data_Factors.parquet'))
+    return ff3
+    
+    
+def set_blacklist(keyword: list | str, checkpoint_dir:list):
+    """Add one or more keywords to the blacklist file.
+
+    Args:
+        keyword: A single keyword or a list of keywords to add.
+        checkpoint_dir: Directory where the blacklist file is stored.
+
+    Returns:
+        None.
+
+    Notes:
+        The directory is created automatically if it does not already exist.
+    """
+    if isinstance(keyword, str):
+        keyword = [keyword]
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    existing = load_blacklist(checkpoint_dir)
+    updated = list(existing | set(keyword))
+    
+    blacklist_path = os.path.join(checkpoint_dir, "blacklist.json")
+    with open(blacklist_path, "w") as f:
+        json.dump(updated, f, indent=2)
+    print(f"blacklist updated: add{keyword}, total {len(keyword)}")
+    
+def save_blacklist(tickers: list, checkpoint_dir:str):
+    """Persist a list of tickers into the blacklist file.
+
+    Args:
+        tickers: Tickers to append to the blacklist.
+        checkpoint_dir: Directory where ``blacklist.json`` is stored.
+
+    Returns:
+        None.
+    """
+    path = os.path.join(checkpoint_dir, "blacklist.json")
+    existing = load_blacklist(path)
+    updated = list(existing | set(tickers))
+    with open(path, "w") as f:
+        json.dump(updated,f, indent=2)
+    print(f"Blacklist updated: {updated}")
+
+def data_acquisition(tickers:list, start_date:str, end_date:str, batch_size:int, 
+                    max_retries: int =3, wait: int = 60, checkpoint_dir: str = "tmp/checkpoint") -> pd.DataFrame:
+    """Download historical price and volume data in batches with checkpoints.
+
+    Args:
+        tickers: List of tickers to download.
+        start_date: Start date in ``YYYY-MM-DD`` format.
+        end_date: End date in ``YYYY-MM-DD`` format.
+        batch_size: Number of tickers per download batch.
+        max_retries: Maximum retries for each batch download.
+        wait: Seconds to wait between retry attempts.
+        checkpoint_dir: Directory used to store batch checkpoints and logs.
+
+    Returns:
+        A tuple of ``(close_dataframe, volume_dataframe)``.
+
+    Notes:
+        Each batch is cached to disk so repeated runs can resume from existing
+        checkpoint files.
+    """
+    close=[]
+    volume=[]
+    task_signature = hashlib.md5(f'{sorted(tickers)}_{start_date}_{end_date}'.encode()).hexdigest()[:8]
+    task_checkpoint_dir = os.path.join(checkpoint_dir,task_signature)
+    os.makedirs(task_checkpoint_dir,exist_ok=True)
+    
+
+    def download_batch_with_retry(batch: list, start_date: str, end_date:str, batch_index:int  ,max_retries:int =3, wait: int =60) ->pd.DataFrame | None :
+        checkpoint_path = os.path.join(task_checkpoint_dir, f"batch_{batch_index}.parquet")
+        if os.path.exists(checkpoint_path):
+            print(f"{batch_index} already exist")
+            return pd.read_parquet(checkpoint_path)
+
+        for attempt in range(max_retries):
+            try:
+                data = yf.download(batch, start=start_date, end=end_date,
+                                auto_adjust=True, progress=False)
+                if data.empty:
+                    raise ValueError("Empty data returned")
+                if isinstance(data.columns, pd.MultiIndex):
+                    close = data["Close"]
+                    empty_tickers = close.columns[close.isnull().all()].tolist()
+                    if empty_tickers:
+                        save_blacklist(empty_tickers, task_checkpoint_dir)
+                        print(f"Blacklisted :{empty_tickers}")
+                data.to_parquet(checkpoint_path)
+                return data
+            except Exception as e:
+                print(f"Attempt {attempt+1} failed: {e}")
+                if attempt < max_retries - 1:
+                    print(f"Retrying after {wait} seconds...")
+                    time.sleep(wait)
+        with open(os.path.join(task_checkpoint_dir, "failed_batches.json"), mode= "a") as f:
+            failed_log = os.path.join(task_checkpoint_dir, "failed_batches.json") 
+            existing_failed = set()
+            if os.path.exists(failed_log):#先读已有的记录，去重后再写
+                with open(failed_log) as f:
+                    for line in f:
+                        existing_failed.add(json.loads(line)["batch_index"])
+
+            if batch_index not in existing_failed:
+                with open(failed_log, mode="a") as f:
+                    json.dump({"batch_index": batch_index, "tickers": batch}, f)
+                    f.write("\n")
+        print(f"Batch {batch_index} failed, logging to {os.path.join(checkpoint_dir, 'failed_batches.json')}")    
+        return None    
+    
+    all_close = []
+    all_volume = []
+    for i in range(0, len(tickers), batch_size):
+        batch = tickers[i: i+batch_size]
+        batch_index= i//batch_size+1
+        print(f"downloading batch{batch_index} ... \ntickers {i} - {min(len(tickers), (i+batch_size))}")
+        data = download_batch_with_retry(batch=batch, start_date=start_date, end_date=end_date, max_retries=max_retries, wait=wait, batch_index=batch_index)
+        if data is not None:
+            if isinstance(data.columns, pd.MultiIndex):
+                all_close.append(data['Close'])
+                all_volume.append(data['Volume'])
+        time.sleep(10)
+    
+    close = pd.concat(all_close, axis=1)
+    volume = pd.concat(all_volume, axis=1)
+    return close, volume
+
+def retry_batches(start_date: str, end_date: str, max_retries: int, checkpoint_dir: str = "tmp/checkpoint", wait: int =60)->pd.DataFrame | None:
+    """Retry batches that previously failed to download.
+
+    Args:
+        start_date: Start date in ``YYYY-MM-DD`` format.
+        end_date: End date in ``YYYY-MM-DD`` format.
+        max_retries: Maximum retries per failed batch.
+        checkpoint_dir: Directory containing failure logs and checkpoints.
+        wait: Seconds to wait between retry attempts.
+
+    Returns:
+        A tuple of ``(close_dataframe, volume_dataframe)`` if any batch is
+        recovered, otherwise ``(None, None)``.
+
+    Notes:
+        The function reads ``failed_batches.json`` and retries only the batches
+        recorded there.
+    """
+    black_list = load_blacklist(checkpoint_dir)
+    failed_log = os.path.join(checkpoint_dir, "failed_batches.json") #读取失败日志
+    retry_log = os.path.join(checkpoint_dir, "retry_log.json") #读取重试日志
+    if not os.path.exists(failed_log):
+        print("No failed batches need to be processed")
+        return None, None
+    failed = []
+    with open(failed_log) as f:
+        for line in f:
+            failed.append(json.loads(line))
+
+    retry_count = 1
+    if os.path.exists(retry_log):
+        with open(retry_log) as f: #读取重试日志，获取上一次是第几次重试
+            records = [json.loads(line) for line in f if line.strip()]
+            if records:
+                retry_count = records[-1]["retry_call"] +1
+    print(f"Retry round {retry_count}, {len(failed)} batches failed in total")
+
+    
+    still_failed=[]
+    retry_records = []
+    success_close=[]
+    success_volume=[]
+
+    for record in failed:
+        batch_index = record["batch_index"]
+        batch = record["tickers"]
+        
+        checkpoint_path = os.path.join(checkpoint_dir, f"batch_{batch_index}.parquet")
+        print(f"Retrying batch {batch_index}")
+        success = False
+        for attempt in range(max_retries):
+            try:
+                data = yf.download(batch, start=start_date, end=end_date,
+                                auto_adjust=True, progress=False)
+                if data.empty:
+                    raise ValueError("Empty data returned")
+                retry_records.append({
+                    "retry_call":retry_count,
+                    "batch_index": batch_index,
+                    "status": "success",
+                    "attempts": attempt + 1
+                })
+                data.to_parquet(checkpoint_path)
+                if isinstance(data.columns, pd.MultiIndex):
+                    success_close.append(data['Close'])
+                    success_volume.append(data['Volume'])
+                print(f"{batch} succeeded!")
+                success = True
+                break
+            except Exception as e:
+                print(f"Attempt {attempt+1} failed: {e}")
+                if attempt < max_retries - 1:
+                    print(f"Retrying after {wait} seconds...")
+                    time.sleep(wait)
+        if not success:
+            print(f"Batch {batch} still failed to download")
+            retry_records.append({
+                    "retry_call":retry_count,
+                    "batch_index": batch_index,
+                    "status": "failed",
+                    "attempts": max_retries
+                })
+            still_failed.append(record)
+        
+        if not batch:
+            print(f"Batch {record["batch_index"]} is fully delisted, skipping")
+            retry_records.append({
+                "retry_call": retry_count,
+                "batch_index":record["batch_index"],
+                "status":"skipped_blacklist",
+                "attempts":0,
+            })
+            continue
+        batch = [t for t in record["tickers"] if t not in black_list]
+        
+        
+    with open(failed_log, "w") as f:
+        for record in still_failed: #将仍失败的任务写在失败日志里
+            json.dump(record, f)
+            f.write("\n")
+
+    with open(retry_log, "a") as f: #将重试记录写在重试的日志里
+        for record in retry_records:
+            json.dump(record, f)
+            f.write("\n")
+    if still_failed:
+        print(f"Still {len(still_failed)} batches failed to download")
+    else:
+        print("All previously failed batches succeeded")
+    if success_volume:
+        volume = pd.concat(success_volume,axis=1)
+        close = pd.concat(success_close,axis=1)
+        return close, volume
+    else:
+        return None, None
+    
+
+def merge_checkpoints(checkpoint_dir: str = "tmp/checkpoint") -> tuple:
+    """Merge all batch checkpoint files into full close and volume tables.
+
+    Args:
+        checkpoint_dir: Directory containing ``batch_*.parquet`` checkpoint
+            files.
+
+    Returns:
+        A tuple of ``(close_dataframe, volume_dataframe)``.
+
+    Notes:
+        Only parquet files with a MultiIndex column layout are merged.
+    """
+    files = sorted(glob.glob(os.path.join(checkpoint_dir, "batch_*.parquet")),
+                   key=lambda x: int(x.split("batch_")[1].split(".")[0]))
+    
+    all_close, all_volume = [], []
+    for f in files:
+        data = pd.read_parquet(f)
+        if isinstance(data.columns, pd.MultiIndex):
+            all_close.append(data["Close"])
+            all_volume.append(data["Volume"])
+    
+    close = pd.concat(all_close, axis=1)
+    volume = pd.concat(all_volume, axis=1)
+    return close, volume

quantmine/datareader.py

@@ -0,0 +1,119 @@
+from dataclasses import dataclass
+import pandas as pd
+from typing import TypeAlias, Protocol
+from abc import ABC, abstractmethod
+
+@dataclass
+class MarketData:
+    close: pd.DataFrame = None
+    volume: pd.DataFrame = None
+    market_cap : pd.DataFrame = None
+    def require(self, *fields: str)->None:
+        for field in fields:
+            if getattr(self, field) is None:
+                raise ValueError(f'You have not provided the data for "{field}"')
+            
+class DataSource(Protocol):
+    def load(self, tickers: list[str], start: str, end: str )-> MarketData:
+        ... #ellipsis: 占位符，约定俗成这个函数体是空的，只是展位
+
+class LocalFileSource(ABC):
+    def __init__(self, paths: dict[str]):
+        self.paths = paths
+    @abstractmethod
+    def _read_file(self, path:str)->pd.DataFrame:
+        ...
+    def load(self, tickers: list[str] =None, start: str=None, end:str= None)-> MarketData:
+        loaded = {}
+        for filed_name, path in self.paths.items():
+            df = self._read_file(path)
+            if tickers:
+                loaded[filed_name] = df.loc[start:end, tickers]
+            else: 
+                loaded[filed_name] = df.loc[start:end]
+        return MarketData(**loaded)
+
+class YFinanceSource:
+    def __init__(self, batch_size: int =50, max_retries: int=3, wait: int=60, checkpoint_dir:str="tmp/checkpoint"):
+        self.batch_size = batch_size
+        self.max_retries = max_retries
+        self.wait = wait
+        self.checkpoint_dir = checkpoint_dir
+    def load(self, tickers: list[str], start: str, end: str) ->MarketData:
+        #惰性import: 只用本地文件源的用户不必安装yfinance依赖链
+        from data_acquisition import data_acquisition
+        close, volume = data_acquisition(tickers=tickers, start_date = start, end_date = end, batch_size= self.batch_size,
+                                    max_retries= self.max_retries, wait= self.wait, checkpoint_dir=self.checkpoint_dir)
+        return MarketData(close = close, volume = volume)
+
+class ParquetSource(LocalFileSource):
+    def _read_file(self, path:str) -> pd.DataFrame:
+        return pd.read_parquet(path)
+
+class CSVSource(LocalFileSource):
+    def _read_file(self, path):
+        return pd.read_csv(path, index_col = 0, parse_dates= True)
+
+class ExcelSource(LocalFileSource):
+    def _read_file(self, path):
+        return pd.read_excel(path, index_col=0, parse_dates= True)
+
+class ConstituentsSource(Protocol):
+    """Point-in-time股票池数据源协议。
+
+    回测在每个调仓日调用get_constituents(date)获取当日有效宇宙,
+    任何实现了该方法的对象都可以接入(静态名单/区间表/实时API)。
+    """
+    def get_constituents(self, date: pd.Timestamp) -> set[str]:
+        ...
+
+class StaticUniverse:
+    """固定股票池: 没有point-in-time成分股数据的用户的最简实现。
+
+    注意: 固定宇宙意味着回测带有幸存者偏差, 结论解读时需要声明。
+    """
+    def __init__(self, tickers: list[str] | set[str]):
+        self._tickers = set(tickers)
+    def get_constituents(self, date: pd.Timestamp) -> set[str]:
+        return self._tickers
+
+class MembershipTableSource:
+    """区间表实现: 从 (ticker, start_date, end_date) 表构造point-in-time宇宙。
+
+    end_date为空表示至今仍在指数内。日期解析在构造时一次完成
+    (每次查询重复parse是逐调仓日的全表开销), 查询结果按日期缓存。
+
+    Args:
+        table: 含ticker/start_date/end_date三列的DataFrame。
+        normalize: True时把ticker里的'.'替换成'-'(如BRK.B -> BRK-B),
+            与yfinance列名对齐; 若你的行情数据用'.'写法则传False。
+    """
+    def __init__(self, table: pd.DataFrame, ticker_col: str = 'ticker',
+                 start_col: str = 'start_date', end_col: str = 'end_date',
+                 normalize: bool = True):
+        self._start = pd.to_datetime(table[start_col])
+        self._end = pd.to_datetime(table[end_col])
+        tickers = table[ticker_col].astype(str)
+        if normalize:
+            tickers = tickers.str.replace('.', '-', regex=False)
+        self._tickers = tickers
+        self._cache: dict[pd.Timestamp, set[str]] = {}
+
+    def get_constituents(self, date: pd.Timestamp) -> set[str]:
+        date = pd.Timestamp(date)
+        if date not in self._cache:
+            #end条件必须整体括起来, 否则end>=date的行会绕过start<=date的检查
+            mask = (self._start <= date) & (self._end.isnull() | (self._end >= date))
+            self._cache[date] = set(self._tickers[mask])
+        return self._cache[date]
+
+FactorDict: TypeAlias = dict[str, pd.DataFrame]
+ForwardReturn: TypeAlias = dict[int, pd.DataFrame]
+
+
+
+
+if __name__ == "__main__":
+    yfinance_data = YFinanceSource()
+    apple_data = yfinance_data.load(tickers = ['AAPL'],start = '2026-07-02', end = '2026-07-08')
+    print(apple_data)

quantmine/factor_attribution.py

@@ -0,0 +1,70 @@
+"""Carhart four-factor attribution of the long-short backtest returns.
+
+Regresses the daily long-short return series on the Fama-French three factors
+plus momentum (all taken from the Ken French Data Library) with Newey-West
+(HAC) standard errors. This answers whether the strategy return is explained
+by known risk premia or carries unexplained alpha.
+
+Factor data (not included in the repo, download from the Ken French Data
+Library and place under ``tmp/ff3/``):
+    - F-F_Research_Data_Factors_daily.csv
+    - F-F_Momentum_Factor_daily.csv
+"""
+import pandas as pd
+import statsmodels.api as sm
+
+
+def load_french_factors(ff3_path: str, mom_path: str) -> pd.DataFrame:
+    """Load and merge daily FF3 and momentum factors from Ken French CSVs.
+
+    Args:
+        ff3_path: Path to ``F-F_Research_Data_Factors_daily.csv``.
+        mom_path: Path to ``F-F_Momentum_Factor_daily.csv``.
+
+    Returns:
+        A dataframe indexed by date with columns ``Mkt-RF``, ``SMB``, ``HML``,
+        ``RF`` and ``Mom``, converted from percent to decimal returns.
+
+    Notes:
+        The raw CSVs carry header/footer junk rows; only rows whose date field
+        matches ``YYYYMMDD`` are kept.
+    """
+    ff3 = pd.read_csv(ff3_path, skiprows=4)
+    ff3.columns = ['Date'] + list(ff3.columns[1:])
+    ff3 = ff3[ff3['Date'].astype(str).str.match(r'^\d{8}$')]
+
+    mom = pd.read_csv(mom_path, skiprows=13)
+    mom = mom.iloc[:, :2]
+    mom.columns = ['Date', 'Mom']
+    mom = mom[mom['Date'].astype(str).str.match(r'^\d{8}$')]
+
+    merged = ff3.merge(mom, on='Date', how='inner')
+    merged['date'] = pd.to_datetime(merged['Date'], format='%Y%m%d')
+    cols = ['Mkt-RF', 'SMB', 'HML', 'RF', 'Mom']
+    merged[cols] = merged[cols].astype(float) / 100  #percent -> decimal
+    return merged.set_index('date')[cols]
+
+
+def carhart_attribution(daily_returns: pd.DataFrame, factors: pd.DataFrame, maxlags: int = 20):
+    """Regress the long-short series on Mkt-RF, SMB, HML and Mom.
+
+    Args:
+        daily_returns: Wide dataframe containing a ``long_short`` column of
+            daily returns (output of ``expand_to_daily_returns``).
+        factors: Daily factor dataframe from ``load_french_factors``.
+        maxlags: Newey-West lag length for the HAC covariance. Default 20 to
+            match the longest holding period.
+
+    Returns:
+        A fitted statsmodels OLS results object. ``params['const']`` is the
+        daily alpha; multiply by 252 to annualize.
+
+    Notes:
+        The long-short portfolio is self-financing, so the raw spread (not the
+        excess over RF) is regressed.
+    """
+    combined = (daily_returns[['long_short']]
+                .join(factors[['Mkt-RF', 'SMB', 'HML', 'Mom']], how='inner')
+                .dropna())
+    x = sm.add_constant(combined[['Mkt-RF', 'SMB', 'HML', 'Mom']])
+    return sm.OLS(combined['long_short'], x).fit(cov_type='HAC', cov_kwds={'maxlags': maxlags})