79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
"""合成数据（MVP 零依赖可跑）。
|
||
|
||
生成随机游走 OHLCV，灌入 DB。仅用于让回测链路在没有真实数据源时也能跑通演示。
|
||
阶段1 接 Tushare/AKShare 后，这里仅保留为"离线测试夹具"。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
import numpy as np
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from ..domain import Bar
|
||
from ..models import Candle
|
||
|
||
|
||
def _trading_days(n: int) -> list[datetime]:
|
||
"""粗略生成 n 个工作日（跳过周末；节假日由阶段1 的交易日历服务处理）。"""
|
||
start = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(days=int(n * 1.6))
|
||
days: list[datetime] = []
|
||
d = start
|
||
while len(days) < n:
|
||
if d.weekday() < 5:
|
||
days.append(d.replace(hour=15, minute=0, second=0, microsecond=0))
|
||
d += timedelta(days=1)
|
||
return days
|
||
|
||
|
||
def generate_ohlcv(n: int = 500, seed: int = 42) -> list[Bar]:
|
||
"""随机游走 + A 股风格的价格区间（5~30 元）。"""
|
||
rng = np.random.default_rng(seed)
|
||
rets = rng.normal(loc=0.0003, scale=0.018, size=n)
|
||
price = 10.0 * np.cumprod(1 + rets)
|
||
days = _trading_days(n)
|
||
bars: list[Bar] = []
|
||
for i in range(n):
|
||
close = float(price[i])
|
||
op = close * (1 + rng.normal(0, 0.005))
|
||
hi = max(op, close) * (1 + abs(rng.normal(0, 0.006)))
|
||
lo = min(op, close) * (1 - abs(rng.normal(0, 0.006)))
|
||
vol = float(rng.integers(1_000_000, 10_000_000))
|
||
bars.append(
|
||
Bar(
|
||
ts=days[i],
|
||
open=round(op, 2),
|
||
high=round(hi, 2),
|
||
low=round(lo, 2),
|
||
close=round(close, 2),
|
||
volume=vol,
|
||
)
|
||
)
|
||
return bars
|
||
|
||
|
||
async def seed_if_empty(session: AsyncSession, symbol: str = "DEMO", n: int = 500) -> None:
|
||
"""若库中无该 symbol 数据，则灌入合成数据。"""
|
||
existing = await session.execute(
|
||
select(Candle.id).where(Candle.symbol == symbol).limit(1)
|
||
)
|
||
if existing.scalars().first() is not None:
|
||
return
|
||
|
||
bars = generate_ohlcv(n=n)
|
||
for b in bars:
|
||
session.add(
|
||
Candle(
|
||
symbol=symbol,
|
||
timeframe="1d",
|
||
ts=b.ts,
|
||
open=b.open,
|
||
high=b.high,
|
||
low=b.low,
|
||
close=b.close,
|
||
volume=b.volume,
|
||
)
|
||
)
|
||
await session.commit()
|