Add minute bar feature pipeline
This commit is contained in:
Yuxuan Yan
@@ -347,3 +347,69 @@ def test_universe_filter_does_not_corrupt_signal_history():
|
||||
held = set(filtered.loc[filtered["weight"] != 0.0, "symbol_id"].unique())
|
||||
# The two most liquid names (highest amount) are sh600519, sz300750.
|
||||
assert held == {"sh600519", "sz300750"}
|
||||
|
||||
|
||||
# --- feature-aware alpha integration ----------------------------------------
|
||||
|
||||
def test_compute_alpha_without_feature_path_matches_empty_feature_paths():
|
||||
data = _make_data()
|
||||
|
||||
base = compute_alpha(data, "rev5", "reversal", lookback=5)
|
||||
with_empty_features = compute_alpha(
|
||||
data,
|
||||
"rev5",
|
||||
"reversal",
|
||||
lookback=5,
|
||||
feature_paths=[],
|
||||
)
|
||||
|
||||
pd.testing.assert_frame_equal(base, with_empty_features)
|
||||
|
||||
|
||||
def test_feature_aware_alpha_reads_joined_feature_column(tmp_path):
|
||||
module_path = tmp_path / "feature_aware_alpha.py"
|
||||
module_path.write_text(textwrap.dedent('''
|
||||
import pandas as pd
|
||||
from pipeline.alpha.base import BaseAlpha
|
||||
from pipeline.alpha.registry import register_alpha
|
||||
|
||||
@register_alpha
|
||||
class FeatureAwareAlpha(BaseAlpha):
|
||||
name = "feature_aware_test_alpha"
|
||||
|
||||
def signal_from_data(
|
||||
self,
|
||||
data: pd.DataFrame,
|
||||
close: pd.DataFrame,
|
||||
) -> pd.DataFrame:
|
||||
signal = data.pivot_table(
|
||||
index="date",
|
||||
columns="symbol_id",
|
||||
values="toy_feature",
|
||||
aggfunc="first",
|
||||
)
|
||||
return signal.reindex(index=close.index, columns=close.columns)
|
||||
'''))
|
||||
|
||||
data = _make_data()
|
||||
feature = data[["symbol_id", "date"]].copy()
|
||||
feature["toy_feature"] = feature["symbol_id"].map({
|
||||
"sh600000": 1.0,
|
||||
"sz000001": 2.0,
|
||||
"sh600519": 3.0,
|
||||
})
|
||||
feature_path = tmp_path / "toy_feature.pq"
|
||||
feature.to_parquet(feature_path, index=False)
|
||||
|
||||
load_alpha_module(str(module_path))
|
||||
result = compute_alpha(
|
||||
data,
|
||||
"feature_run",
|
||||
"feature_aware_test_alpha",
|
||||
feature_paths=[str(feature_path)],
|
||||
)
|
||||
|
||||
assert list(result.columns) == ALPHA_COLUMNS
|
||||
assert (result["alpha_name"] == "feature_run").all()
|
||||
last = result[result["date"] == result["date"].max()]
|
||||
assert last.set_index("symbol_id")["weight"].idxmax() == "sh600519"
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Tests for minute-derived daily feature plugins."""
|
||||
|
||||
import textwrap
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from pipeline.features.compute import compute_feature, validate_feature_frame
|
||||
from pipeline.features.registry import (
|
||||
available_features,
|
||||
get_feature,
|
||||
load_feature_module,
|
||||
)
|
||||
|
||||
|
||||
def _minute_bars() -> pd.DataFrame:
|
||||
return pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sh600000", "sz000001"],
|
||||
"symbol_name": ["PF Bank", "PF Bank", "Ping An"],
|
||||
"datetime": pd.to_datetime([
|
||||
"2024-01-02 09:35:00",
|
||||
"2024-01-02 09:40:00",
|
||||
"2024-01-02 09:35:00",
|
||||
]),
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-02"]),
|
||||
"time": ["09:35:00", "09:40:00", "09:35:00"],
|
||||
"frequency": ["5m", "5m", "5m"],
|
||||
"open": [10.0, 10.5, 20.0],
|
||||
"high": [11.0, 12.0, 21.0],
|
||||
"low": [9.0, 10.0, 19.0],
|
||||
"close": [10.5, 11.0, 20.5],
|
||||
"volume": [100.0, 300.0, 200.0],
|
||||
"amount": [1000.0, 3300.0, 4100.0],
|
||||
"vwap": [10.0, 11.0, 20.5],
|
||||
"adjustflag": ["3", "3", "3"],
|
||||
})
|
||||
|
||||
|
||||
def test_built_in_minute_daily_summary():
|
||||
daily = pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sz000001", "sh600000"],
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-03"]),
|
||||
"close": [11.0, 20.5, 12.0],
|
||||
})
|
||||
|
||||
result = compute_feature(
|
||||
minute=_minute_bars(),
|
||||
daily=daily,
|
||||
feature_type="minute_daily_summary",
|
||||
)
|
||||
|
||||
assert "minute_daily_summary" in available_features()
|
||||
row = result[
|
||||
(result["symbol_id"] == "sh600000")
|
||||
& (result["date"] == pd.Timestamp("2024-01-02"))
|
||||
].iloc[0]
|
||||
assert row["minute_bar_count"] == 2
|
||||
assert np.isclose(row["minute_intraday_return"], 11.0 / 10.0 - 1.0)
|
||||
assert np.isclose(row["minute_intraday_range"], 12.0 / 9.0 - 1.0)
|
||||
assert np.isclose(row["minute_vwap"], 4300.0 / 400.0)
|
||||
assert np.isclose(row["minute_vwap_deviation"], (4300.0 / 400.0) / 11.0 - 1.0)
|
||||
|
||||
missing = result[
|
||||
(result["symbol_id"] == "sh600000")
|
||||
& (result["date"] == pd.Timestamp("2024-01-03"))
|
||||
].iloc[0]
|
||||
assert pd.isna(missing["minute_vwap"])
|
||||
|
||||
|
||||
def test_load_external_feature_module_and_filter_params(tmp_path):
|
||||
module_path = tmp_path / "external_feature.py"
|
||||
module_path.write_text(textwrap.dedent('''
|
||||
import pandas as pd
|
||||
from pipeline.features.base import BaseFeature
|
||||
from pipeline.features.registry import register_feature
|
||||
|
||||
@register_feature
|
||||
class ExternalVolumeFeature(BaseFeature):
|
||||
name = "external_volume_feature"
|
||||
|
||||
def __init__(self, scale: float = 1.0):
|
||||
self.scale = scale
|
||||
|
||||
def compute(self, minute: pd.DataFrame, daily=None) -> pd.DataFrame:
|
||||
out = (
|
||||
minute.groupby(["symbol_id", "date"], as_index=False)["volume"]
|
||||
.sum()
|
||||
.rename(columns={"volume": "scaled_volume"})
|
||||
)
|
||||
out["scaled_volume"] *= self.scale
|
||||
return out
|
||||
'''))
|
||||
|
||||
load_feature_module(str(module_path))
|
||||
assert "external_volume_feature" in available_features()
|
||||
|
||||
instance = get_feature("external_volume_feature", scale=2.0, ignored=99)
|
||||
assert instance.scale == 2.0
|
||||
assert not hasattr(instance, "ignored")
|
||||
|
||||
result = compute_feature(
|
||||
minute=_minute_bars(),
|
||||
feature_type="external_volume_feature",
|
||||
scale=2.0,
|
||||
ignored=99,
|
||||
)
|
||||
row = result[result["symbol_id"] == "sh600000"].iloc[0]
|
||||
assert np.isclose(row["scaled_volume"], 800.0)
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_missing_keys():
|
||||
with pytest.raises(ValueError, match="missing required"):
|
||||
validate_feature_frame(pd.DataFrame({"symbol_id": ["sh600000"], "x": [1.0]}))
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_duplicate_keys_after_date_normalization():
|
||||
with pytest.raises(ValueError, match="duplicate symbol_id,date"):
|
||||
validate_feature_frame(pd.DataFrame({
|
||||
"symbol_id": ["sh600000", "sh600000"],
|
||||
"date": ["2024-01-02", pd.Timestamp("2024-01-02")],
|
||||
"x": [1.0, 2.0],
|
||||
}))
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_duplicate_columns():
|
||||
bad = pd.DataFrame(
|
||||
[["sh600000", pd.Timestamp("2024-01-02"), 1.0, 2.0]],
|
||||
columns=["symbol_id", "date", "dup", "dup"],
|
||||
)
|
||||
with pytest.raises(ValueError, match="duplicate columns"):
|
||||
validate_feature_frame(bad)
|
||||
|
||||
|
||||
def test_validate_feature_frame_rejects_non_numeric_feature_columns():
|
||||
with pytest.raises(ValueError, match="numeric"):
|
||||
validate_feature_frame(pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"date": [pd.Timestamp("2024-01-02")],
|
||||
"bad": ["not numeric"],
|
||||
}))
|
||||
@@ -0,0 +1,215 @@
|
||||
"""Tests for raw Baostock minute bar download plumbing."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import data.downloader as low_level_downloader
|
||||
import pipeline.data.downloader as pipeline_downloader
|
||||
from data.downloader import download_minute_batch
|
||||
from pipeline.common.schema import MINUTE_BAR_COLUMNS
|
||||
from pipeline.data.downloader import download_minute_universe
|
||||
|
||||
|
||||
class _FakeResult:
|
||||
def __init__(self, rows, error_code="0", error_msg=""):
|
||||
self.rows = rows
|
||||
self.error_code = error_code
|
||||
self.error_msg = error_msg
|
||||
self._idx = -1
|
||||
|
||||
def next(self):
|
||||
self._idx += 1
|
||||
return self._idx < len(self.rows)
|
||||
|
||||
def get_row_data(self):
|
||||
return self.rows[self._idx]
|
||||
|
||||
|
||||
def test_download_minute_batch_maps_and_parses_baostock_rows(monkeypatch):
|
||||
rows = [
|
||||
[
|
||||
"2024-01-02",
|
||||
"20240102093500000",
|
||||
"sh.600000",
|
||||
"10",
|
||||
"11",
|
||||
"9",
|
||||
"10.5",
|
||||
"1000",
|
||||
"10500",
|
||||
"3",
|
||||
],
|
||||
[
|
||||
"2024-01-02",
|
||||
"20240102094000000",
|
||||
"sh.600000",
|
||||
"10.5",
|
||||
"12",
|
||||
"10",
|
||||
"11",
|
||||
"2000",
|
||||
"22000",
|
||||
"3",
|
||||
],
|
||||
]
|
||||
calls = []
|
||||
|
||||
def fake_query(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return _FakeResult(rows)
|
||||
|
||||
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
|
||||
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
low_level_downloader.bs,
|
||||
"query_history_k_data_plus",
|
||||
fake_query,
|
||||
)
|
||||
|
||||
[(symbol, df)] = list(
|
||||
download_minute_batch(
|
||||
["sh600000"],
|
||||
"2024-01-02",
|
||||
"2024-01-02",
|
||||
frequency=5,
|
||||
)
|
||||
)
|
||||
|
||||
assert symbol == "sh600000"
|
||||
assert calls[0]["code"] == "sh.600000"
|
||||
assert calls[0]["frequency"] == "5"
|
||||
assert calls[0]["adjustflag"] == "3"
|
||||
assert df is not None
|
||||
assert df["datetime"].iloc[0] == pd.Timestamp("2024-01-02 09:35:00")
|
||||
assert df["time"].tolist() == ["09:35:00", "09:40:00"]
|
||||
assert (df["frequency"] == "5m").all()
|
||||
assert np.isclose(df["open"].iloc[0], 10.0)
|
||||
assert np.isclose(df["vwap"].iloc[0], 10.5)
|
||||
assert pd.api.types.is_numeric_dtype(df["volume"])
|
||||
|
||||
|
||||
def test_download_minute_batch_empty_result_yields_none(monkeypatch):
|
||||
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
|
||||
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
low_level_downloader.bs,
|
||||
"query_history_k_data_plus",
|
||||
lambda **kwargs: _FakeResult([]),
|
||||
)
|
||||
|
||||
assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
|
||||
("sh600000", None)
|
||||
]
|
||||
|
||||
|
||||
def test_download_minute_batch_rejects_unparsed_timestamps(monkeypatch):
|
||||
bad_rows = [[
|
||||
"2024-01-02",
|
||||
"not-a-time",
|
||||
"sh.600000",
|
||||
"10",
|
||||
"11",
|
||||
"9",
|
||||
"10.5",
|
||||
"1000",
|
||||
"10500",
|
||||
"3",
|
||||
]]
|
||||
monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None)
|
||||
monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
low_level_downloader.bs,
|
||||
"query_history_k_data_plus",
|
||||
lambda **kwargs: _FakeResult(bad_rows),
|
||||
)
|
||||
|
||||
assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [
|
||||
("sh600000", None)
|
||||
]
|
||||
|
||||
|
||||
def test_download_minute_universe_writes_frequency_month_partitions(tmp_path, monkeypatch):
|
||||
minute = pd.DataFrame({
|
||||
"symbol": ["sh600000", "sh600000"],
|
||||
"datetime": pd.to_datetime(["2024-01-02 09:35:00", "2024-01-02 09:40:00"]),
|
||||
"date": pd.to_datetime(["2024-01-02", "2024-01-02"]),
|
||||
"time": ["09:35:00", "09:40:00"],
|
||||
"frequency": ["5m", "5m"],
|
||||
"open": [10.0, 10.5],
|
||||
"high": [11.0, 12.0],
|
||||
"low": [9.0, 10.0],
|
||||
"close": [10.5, 11.0],
|
||||
"volume": [1000.0, 2000.0],
|
||||
"amount": [10500.0, 22000.0],
|
||||
"vwap": [10.5, 11.0],
|
||||
"adjustflag": ["3", "3"],
|
||||
})
|
||||
|
||||
monkeypatch.setattr(
|
||||
pipeline_downloader,
|
||||
"_resolve_universe",
|
||||
lambda universe, max_symbols=0: pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"symbol_name": ["PF Bank"],
|
||||
}),
|
||||
)
|
||||
|
||||
def fake_batch(symbols, start, end, frequency=5):
|
||||
assert symbols == ["sh600000"]
|
||||
assert frequency == "5"
|
||||
yield "sh600000", minute
|
||||
|
||||
monkeypatch.setattr(pipeline_downloader, "download_minute_batch", fake_batch)
|
||||
|
||||
preserved = tmp_path / "toy" / "frequency=15m" / "month=2024-01" / "old.pq"
|
||||
preserved.parent.mkdir(parents=True)
|
||||
preserved_minute = minute.copy()
|
||||
preserved_minute["frequency"] = "15m"
|
||||
preserved_minute["symbol_id"] = "sh600000"
|
||||
preserved_minute["symbol_name"] = "PF Bank"
|
||||
preserved_minute[MINUTE_BAR_COLUMNS].to_parquet(preserved, index=False)
|
||||
|
||||
stats = download_minute_universe(
|
||||
universe="toy",
|
||||
start_date="2024-01-02",
|
||||
end_date="2024-01-02",
|
||||
output_dir=str(tmp_path),
|
||||
chunk_size=1,
|
||||
frequency="5",
|
||||
)
|
||||
|
||||
dataset_path = Path(stats["dataset_path"])
|
||||
assert (dataset_path / "frequency=5m" / "month=2024-01").is_dir()
|
||||
assert preserved.exists()
|
||||
out = pd.read_parquet(dataset_path / "frequency=5m")
|
||||
assert (set(MINUTE_BAR_COLUMNS) - {"frequency"}) <= set(out.columns)
|
||||
assert set(out["symbol_id"]) == {"sh600000"}
|
||||
assert set(out["symbol_name"]) == {"PF Bank"}
|
||||
assert stats["n_rows"] == 2
|
||||
|
||||
|
||||
def test_download_minute_universe_raises_when_all_symbols_empty(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
pipeline_downloader,
|
||||
"_resolve_universe",
|
||||
lambda universe, max_symbols=0: pd.DataFrame({
|
||||
"symbol_id": ["sh600000"],
|
||||
"symbol_name": ["PF Bank"],
|
||||
}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
pipeline_downloader,
|
||||
"download_minute_batch",
|
||||
lambda symbols, start, end, frequency=5: iter([("sh600000", None)]),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="No minute data"):
|
||||
download_minute_universe(
|
||||
universe="toy",
|
||||
start_date="2024-01-02",
|
||||
end_date="2024-01-02",
|
||||
output_dir=str(tmp_path),
|
||||
)
|
||||
Reference in New Issue
Block a user