Add JoinQuant comparison plugin
This commit is contained in:
Yuxuan Yan
@@ -0,0 +1,201 @@
|
||||
"""Export portfolio positions as frozen JoinQuant target files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Literal
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from pipeline.common.schema import POSITION_COLUMNS
|
||||
from plugins.joinquant.schema import JOINQUANT_TARGET_COLUMNS
|
||||
from plugins.joinquant.symbols import to_joinquant_symbol
|
||||
|
||||
|
||||
ExportMode = Literal["target_shares", "target_value"]
|
||||
|
||||
|
||||
def _date_text(value: object) -> str:
|
||||
return pd.Timestamp(value).strftime("%Y-%m-%d")
|
||||
|
||||
|
||||
def _date_file_stem(date_text: str) -> str:
|
||||
return pd.Timestamp(date_text).strftime("%Y%m%d")
|
||||
|
||||
|
||||
def _snapshot_root_for(targets_root: Path) -> Path:
|
||||
if targets_root.name == "targets":
|
||||
return targets_root.parent / "snapshots"
|
||||
return targets_root / "snapshots"
|
||||
|
||||
|
||||
def _sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as fh:
|
||||
for chunk in iter(lambda: fh.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _check_position_columns(df: pd.DataFrame) -> None:
|
||||
missing = [col for col in POSITION_COLUMNS if col not in df.columns]
|
||||
if missing:
|
||||
raise ValueError(f"Positions input missing required columns: {missing}")
|
||||
|
||||
|
||||
def _filter_dates(
|
||||
df: pd.DataFrame,
|
||||
start_date: str | None,
|
||||
end_date: str | None,
|
||||
) -> pd.DataFrame:
|
||||
out = df.copy()
|
||||
out["date"] = pd.to_datetime(out["date"]).dt.normalize()
|
||||
if start_date:
|
||||
out = out[out["date"] >= pd.Timestamp(start_date).normalize()]
|
||||
if end_date:
|
||||
out = out[out["date"] <= pd.Timestamp(end_date).normalize()]
|
||||
return out
|
||||
|
||||
|
||||
def build_target_frame(
|
||||
positions: pd.DataFrame,
|
||||
*,
|
||||
portfolio_name: str | None = None,
|
||||
mode: ExportMode = "target_shares",
|
||||
start_date: str | None = None,
|
||||
end_date: str | None = None,
|
||||
snapshot_ids: dict[str, str] | None = None,
|
||||
) -> pd.DataFrame:
|
||||
"""Build normalized JoinQuant target rows from portfolio positions.
|
||||
|
||||
``target_shares`` is populated from ``position_shares`` because the core
|
||||
simulator executes the discretized book, not continuous research shares.
|
||||
"""
|
||||
if mode not in {"target_shares", "target_value"}:
|
||||
raise ValueError("mode must be 'target_shares' or 'target_value'")
|
||||
_check_position_columns(positions)
|
||||
|
||||
df = _filter_dates(positions, start_date, end_date)
|
||||
if portfolio_name is not None:
|
||||
df = df[df["portfolio_name"].astype(str) == portfolio_name]
|
||||
if df.empty:
|
||||
return pd.DataFrame(columns=JOINQUANT_TARGET_COLUMNS)
|
||||
|
||||
out = pd.DataFrame({
|
||||
"date": df["date"].map(_date_text),
|
||||
"portfolio_name": df["portfolio_name"].astype(str),
|
||||
"symbol_id": df["symbol_id"].astype(str),
|
||||
"jq_symbol": df["symbol_id"].map(to_joinquant_symbol),
|
||||
"target_shares": pd.to_numeric(df["position_shares"], errors="coerce").fillna(0).astype("int64"),
|
||||
"target_value": pd.to_numeric(df["target_value"], errors="coerce").fillna(0.0),
|
||||
"target_weight": pd.to_numeric(df["target_weight"], errors="coerce").fillna(0.0),
|
||||
"export_mode": mode,
|
||||
"snapshot_id": "",
|
||||
})
|
||||
|
||||
if snapshot_ids:
|
||||
out["snapshot_id"] = out["date"].map(snapshot_ids).fillna("")
|
||||
|
||||
return out[JOINQUANT_TARGET_COLUMNS].sort_values(
|
||||
["date", "portfolio_name", "symbol_id"]
|
||||
).reset_index(drop=True)
|
||||
|
||||
|
||||
def export_targets(
|
||||
positions_path: str | Path,
|
||||
*,
|
||||
portfolio_name: str,
|
||||
mode: ExportMode = "target_shares",
|
||||
out_dir: str | Path = "plugins_output/joinquant/targets",
|
||||
start_date: str | None = None,
|
||||
end_date: str | None = None,
|
||||
force: bool = False,
|
||||
) -> list[dict[str, object]]:
|
||||
"""Export one daily CSV/parquet target file plus a snapshot JSON per date.
|
||||
|
||||
Args:
|
||||
positions_path: Parquet file produced by ``portfolio build``.
|
||||
portfolio_name: Portfolio run to export.
|
||||
mode: ``target_shares`` or ``target_value``.
|
||||
out_dir: Target root. Files are written to ``out_dir/portfolio_name``.
|
||||
If the root is named ``targets``, snapshots are written to the
|
||||
sibling ``snapshots`` directory.
|
||||
start_date: Optional inclusive start date.
|
||||
end_date: Optional inclusive end date.
|
||||
force: If false, existing target or snapshot files are treated as
|
||||
frozen and cause ``FileExistsError``.
|
||||
|
||||
Returns:
|
||||
Snapshot metadata dictionaries, one per exported date.
|
||||
"""
|
||||
positions_path = Path(positions_path)
|
||||
targets_root = Path(out_dir)
|
||||
snapshot_root = _snapshot_root_for(targets_root)
|
||||
targets_portfolio_dir = targets_root / portfolio_name
|
||||
snapshots_portfolio_dir = snapshot_root / portfolio_name
|
||||
targets_portfolio_dir.mkdir(parents=True, exist_ok=True)
|
||||
snapshots_portfolio_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
positions = pd.read_parquet(positions_path)
|
||||
filtered = _filter_dates(positions, start_date, end_date)
|
||||
filtered = filtered[filtered["portfolio_name"].astype(str) == portfolio_name]
|
||||
if filtered.empty:
|
||||
return []
|
||||
|
||||
date_texts = sorted(filtered["date"].map(_date_text).unique())
|
||||
snapshot_ids = {
|
||||
date_text: f"jq-{portfolio_name}-{date_text}-{uuid.uuid4().hex[:12]}"
|
||||
for date_text in date_texts
|
||||
}
|
||||
targets = build_target_frame(
|
||||
filtered,
|
||||
portfolio_name=portfolio_name,
|
||||
mode=mode,
|
||||
snapshot_ids=snapshot_ids,
|
||||
)
|
||||
|
||||
snapshots: list[dict[str, object]] = []
|
||||
for date_text, daily in targets.groupby("date", sort=True):
|
||||
stem = _date_file_stem(date_text)
|
||||
csv_path = targets_portfolio_dir / f"{stem}.csv"
|
||||
parquet_path = targets_portfolio_dir / f"{stem}.parquet"
|
||||
snapshot_path = snapshots_portfolio_dir / f"{stem}.json"
|
||||
existing: Iterable[Path] = (csv_path, parquet_path, snapshot_path)
|
||||
if not force:
|
||||
conflicts = [str(path) for path in existing if path.exists()]
|
||||
if conflicts:
|
||||
raise FileExistsError(
|
||||
"Frozen JoinQuant target already exists; use --force to overwrite: "
|
||||
+ ", ".join(conflicts)
|
||||
)
|
||||
|
||||
daily = daily[JOINQUANT_TARGET_COLUMNS].reset_index(drop=True)
|
||||
daily.to_csv(csv_path, index=False)
|
||||
daily.to_parquet(parquet_path, index=False)
|
||||
file_hash = _sha256_file(csv_path)
|
||||
|
||||
snapshot = {
|
||||
"snapshot_id": snapshot_ids[date_text],
|
||||
"portfolio_name": portfolio_name,
|
||||
"date": date_text,
|
||||
"export_mode": mode,
|
||||
"source_positions_path": str(positions_path),
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"n_symbols": int(len(daily)),
|
||||
"file_sha256": file_hash,
|
||||
"notes": "Frozen JoinQuant target file.",
|
||||
"target_csv_path": str(csv_path),
|
||||
"target_parquet_path": str(parquet_path),
|
||||
}
|
||||
snapshot_path.write_text(
|
||||
json.dumps(snapshot, indent=2, ensure_ascii=False) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
snapshots.append(snapshot)
|
||||
|
||||
return snapshots
|
||||
|
||||
Reference in New Issue
Block a user