Files
stock/backend/app/data/aggregation.py
2026-08-07 16:08:34 +08:00

55 lines
1.7 KiB
Python

"""K 线周期聚合:日线 -> 周/月/年。
生产环境用 TimescaleDB Continuous Aggregates 在库里预物化(性能);
MVP 在应用层用 pandas resample 即可,逻辑等价、便于切换。
OHLCV 聚合规则:开=周期内首根开、高=最高、低=最低、收=末根收、量=求和。
"""
from __future__ import annotations
import pandas as pd
from ..domain import Bar
# pandas resample 规则(周一为周首;月/年以首日对齐)
_RULES = {"1w": "W-MON", "1M": "MS", "1y": "YS"}
# 各周期的"年交易日数"(用于夏普等指标的年化)
_BARS_PER_YEAR = {"1d": 252, "1w": 52, "1M": 12, "1y": 1}
def bars_per_year(timeframe: str) -> int:
return _BARS_PER_YEAR.get(timeframe, 252)
def resample_bars(bars: list[Bar], timeframe: str) -> list[Bar]:
"""把日线 bars 聚合为目标周期;日线或未知周期原样返回。"""
if not bars or timeframe in ("1d", "d", "day", "", None):
return bars
rule = _RULES.get(timeframe)
if rule is None:
return bars
df = pd.DataFrame(
[{"ts": b.ts, "open": b.open, "high": b.high, "low": b.low, "close": b.close, "volume": b.volume}
for b in bars]
).set_index("ts").sort_index()
agg = (
df.resample(rule)
.agg({"open": "first", "high": "max", "low": "min", "close": "last", "volume": "sum"})
.dropna()
)
return [
Bar(
ts=ts.to_pydatetime(),
open=float(row["open"]),
high=float(row["high"]),
low=float(row["low"]),
close=float(row["close"]),
volume=float(row["volume"]),
)
for ts, row in agg.iterrows()
]