58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import math
|
|
from datetime import datetime
|
|
|
|
import polars as pl
|
|
|
|
from vnpy.alpha.dataset.utility import calculate_by_expression
|
|
|
|
|
|
def sample_frame() -> pl.DataFrame:
|
|
dates: list[datetime] = [
|
|
datetime(2024, 1, 1),
|
|
datetime(2024, 1, 2),
|
|
datetime(2024, 1, 3),
|
|
]
|
|
return pl.DataFrame(
|
|
{
|
|
"datetime": dates + dates,
|
|
"vt_symbol": ["A", "A", "A", "B", "B", "B"],
|
|
"close": [1.0, 3.0, 5.0, 4.0, 2.0, 8.0],
|
|
}
|
|
)
|
|
|
|
|
|
def data_values(frame: pl.DataFrame) -> list[float | None]:
|
|
values: list[float | None] = []
|
|
value: object
|
|
for value in frame["data"].to_list():
|
|
if value is None:
|
|
values.append(None)
|
|
else:
|
|
values.append(float(value))
|
|
return values
|
|
|
|
|
|
def assert_values(actual: list[float | None], expected: list[float | None]) -> None:
|
|
assert len(actual) == len(expected)
|
|
got: float | None
|
|
want: float | None
|
|
for got, want in zip(actual, expected, strict=True):
|
|
if want is None:
|
|
assert got is None
|
|
else:
|
|
assert got is not None
|
|
assert math.isclose(got, want)
|
|
|
|
|
|
def test_delay_mean_and_rank_on_tiny_frame() -> None:
|
|
frame: pl.DataFrame = sample_frame()
|
|
|
|
delayed: pl.DataFrame = calculate_by_expression(frame, "ts_delay(close, 1)")
|
|
averaged: pl.DataFrame = calculate_by_expression(frame, "ts_mean(close, 2)")
|
|
ranked: pl.DataFrame = calculate_by_expression(frame, "cs_rank(close)")
|
|
|
|
assert delayed["vt_symbol"].to_list() == ["A", "A", "A", "B", "B", "B"]
|
|
assert_values(data_values(delayed), [None, 1.0, 3.0, None, 4.0, 2.0])
|
|
assert_values(data_values(averaged), [1.0, 2.0, 4.0, 4.0, 3.0, 5.0])
|
|
assert_values(data_values(ranked), [0.5, 1.0, 0.5, 1.0, 0.5, 1.0])
|