예제 소스

rl_data_labels.py

"""
Market 극단 라벨 CSV(cache/market_extreme_labels.csv 등)를
Market.get_data() 결과 df에 병합한다.

``source_bar_offset``: 원본 전체 시계열에서 ``df`` 가 시작하는 bar_index
(예: tail 25000봉만 쓸 때 len(full)-25000).
"""

from __future__ import annotations

import logging
from pathlib import Path

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# RL 상태(_label_feature_vector)에 쓰는 숫자·불리언 컬럼만 병합
_RL_EXT_COLS = (
    "ext_anomaly_score",
    "ext_crash_signal",
    "ext_surge_signal",
    "ext_price_spike",
    "ext_volume_spike",
    "ext_pattern_anomaly",
    "ext_regime_id",
    "ext_is_anomaly",
)


def merge_extreme_labels_from_csv(
    df: pd.DataFrame,
    csv_path: str | Path,
    *,
    source_bar_offset: int = 0,
) -> pd.DataFrame:
    """
    ``ext_*`` 컬럼을 ``df`` 에 붙인다. ``bar_index`` 가 CSV에 있으면
    ``[source_bar_offset, source_bar_offset + len(df))`` 구간만 사용한다.
    """
    path = Path(csv_path)
    if not path.exists():
        logger.warning("극단 라벨 CSV 없음 — RL 상태의 라벨 슬롯은 0으로 둡니다: %s", path)
        return df

    hdr = pd.read_csv(path, nrows=0).columns.tolist()
    ext_cols = [c for c in _RL_EXT_COLS if c in hdr]
    usecols = ext_cols.copy()
    if "bar_index" in hdr:
        usecols = ["bar_index"] + ext_cols
    logger.info("라벨 CSV 로드 중 (컬럼 %s개) …", len(usecols))
    lab = pd.read_csv(path, usecols=usecols)

    out = df.reset_index(drop=True).copy()
    n = len(out)
    lo = int(source_bar_offset)
    hi = lo + n

    if "bar_index" not in lab.columns:
        lab = lab.copy()
        lab.insert(0, "bar_index", np.arange(lo, lo + len(lab), dtype=np.int64))

    if not ext_cols:
        logger.warning("CSV에 ext_* 컬럼이 없습니다: %s", path)
        return out

    for c in ext_cols:
        if lab[c].dtype == object:
            lab[c] = pd.to_numeric(lab[c], errors="coerce")

    sub = lab.loc[(lab["bar_index"] >= lo) & (lab["bar_index"] < hi)].drop_duplicates(
        subset=["bar_index"], keep="last"
    )
    sub = sub.set_index("bar_index").sort_index().reindex(range(lo, hi))

    if len(sub) != n:
        logger.warning("reindex 후 길이 이상: %s (기대 %s)", len(sub), n)

    for c in ext_cols:
        out[c] = pd.to_numeric(sub[c], errors="coerce").fillna(0.0).to_numpy(dtype=np.float64)

    logger.info("극단 라벨 병합 완료: %s (%s 컬럼, offset=%s)", path, len(ext_cols), source_bar_offset)
    return out