Yuxuan Yan
142 lines
4.7 KiB
Python
142 lines
4.7 KiB
Python
"""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"],
|
|
}))
|