google-labs-jules[bot]
44c4e22e2d
- Added `pytest` to dependencies in `pyproject.toml` - Removed duplicate `create_labels_double_barrier` function definition from `features/labeling_schemes.py` - Added comprehensive unit tests in `tests/test_labeling_schemes.py` covering: - `calculate_future_returns` - `create_labels_multi_bar` - `create_labels_double_barrier` - `create_labels_regime_detection` - `create_labels_volatility` Co-authored-by: maghdam <63883156+maghdam@users.noreply.github.com>
124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
import pytest
|
|
import pandas as pd
|
|
import numpy as np
|
|
from features.labeling_schemes import (
|
|
calculate_future_returns,
|
|
create_labels_multi_bar,
|
|
create_labels_double_barrier,
|
|
create_labels_regime_detection,
|
|
create_labels_volatility,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def sample_data():
|
|
"""Provides a sample DataFrame for testing."""
|
|
# Prices: 100 -> 101 -> 102 -> 100 -> 98 -> 95 -> 100
|
|
dates = pd.date_range("2023-01-01", periods=7, freq="D")
|
|
df = pd.DataFrame({"close": [100.0, 101.0, 102.0, 100.0, 98.0, 95.0, 100.0]}, index=dates)
|
|
return df
|
|
|
|
@pytest.fixture
|
|
def sample_data_long():
|
|
"""Provides a longer sample DataFrame for testing regime detection."""
|
|
# 55 periods
|
|
dates = pd.date_range("2023-01-01", periods=55, freq="D")
|
|
# First 30 periods: trend up, then trend down
|
|
closes = np.linspace(100, 130, 30).tolist() + np.linspace(130, 100, 25).tolist()
|
|
df = pd.DataFrame({"close": closes}, index=dates)
|
|
return df
|
|
|
|
def test_calculate_future_returns(sample_data):
|
|
df_out = calculate_future_returns(sample_data, horizon=1)
|
|
|
|
# Check shape, original is 7, shifted by 1 means 1 nan dropped => 6 rows
|
|
assert len(df_out) == 6
|
|
assert "future_returns" in df_out.columns
|
|
|
|
# 100 to 101 => 0.01
|
|
assert pytest.approx(df_out["future_returns"].iloc[0]) == 0.01
|
|
# 101 to 102 => 1/101
|
|
assert pytest.approx(df_out["future_returns"].iloc[1]) == 1.0 / 101.0
|
|
|
|
def test_create_labels_multi_bar(sample_data):
|
|
df_out = create_labels_multi_bar(sample_data, horizon=2, threshold=0.01)
|
|
|
|
# original is 7, horizon=2 means 2 nan dropped => 5 rows
|
|
assert len(df_out) == 5
|
|
assert "future_return_h" in df_out.columns
|
|
assert "multi_bar_label" in df_out.columns
|
|
|
|
# idx 0: 100 -> 102 (horizon 2), return = 0.02 >= 0.01 => label 1
|
|
assert df_out["multi_bar_label"].iloc[0] == 1
|
|
# idx 2: 102 -> 98 (horizon 2), return = -4/102 = -0.039 <= -0.01 => label -1
|
|
assert df_out["multi_bar_label"].iloc[2] == -1
|
|
# Let's check idx 1: 101 -> 100, return = -1/101 = -0.0099. threshold is 0.01
|
|
# -0.0099 is NOT <= -0.01, so label should be 0
|
|
assert df_out["multi_bar_label"].iloc[1] == 0
|
|
|
|
def test_create_labels_double_barrier(sample_data):
|
|
# Prices: 100 -> 101 -> 102 -> 100 -> 98 -> 95 -> 100
|
|
# up=0.015, down=0.015, horizon=3
|
|
# For idx 0 (100): upper=101.5, lower=98.5
|
|
# horizon=3 => looks at 101, 102, 100
|
|
# hits 102 > 101.5 first. => label 1
|
|
|
|
# For idx 2 (102): upper=103.53, lower=100.47
|
|
# horizon=3 => looks at 100, 98, 95
|
|
# hits 100 < 100.47 first. => label -1
|
|
|
|
# For idx 3 (100): upper=101.5, lower=98.5
|
|
# horizon=3 => looks at 98, 95, 100
|
|
# hits 98 < 98.5 first. => label -1
|
|
|
|
# Let's pick a case where neither is hit:
|
|
# We need a small horizon or large barriers.
|
|
|
|
df_out = create_labels_double_barrier(sample_data, up=0.015, down=0.015, horizon=3)
|
|
|
|
assert "barrier_label" in df_out.columns
|
|
assert len(df_out) == 7
|
|
|
|
assert df_out["barrier_label"].iloc[0] == 1
|
|
assert df_out["barrier_label"].iloc[2] == -1
|
|
assert df_out["barrier_label"].iloc[3] == -1
|
|
|
|
# test no hit
|
|
df_out_nohit = create_labels_double_barrier(sample_data, up=0.1, down=0.1, horizon=1)
|
|
# 100 -> 101 (up=110, down=90). Not hit. => label 0
|
|
assert df_out_nohit["barrier_label"].iloc[0] == 0
|
|
|
|
def test_create_labels_regime_detection(sample_data_long):
|
|
df_out = create_labels_regime_detection(sample_data_long, short_window=5, long_window=10)
|
|
|
|
# dropna for long_window => 55 - 10 + 1 = 46 rows
|
|
assert len(df_out) == 46
|
|
assert "regime_label" in df_out.columns
|
|
assert "ma_short" in df_out.columns
|
|
assert "ma_long" in df_out.columns
|
|
|
|
# Since first 30 are trending up, short MA > long MA
|
|
assert df_out["regime_label"].iloc[0] == 1
|
|
|
|
# The end is trending down, short MA < long MA
|
|
assert df_out["regime_label"].iloc[-1] == -1
|
|
|
|
def test_create_labels_volatility(sample_data):
|
|
# Prices: 100 -> 101 -> 102 -> 100 -> 98 -> 95 -> 100
|
|
# First, calculate returns for horizon=1:
|
|
# 0.01, 0.0099, -0.0196, -0.02, -0.0306, 0.0526
|
|
# Volatility window 2:
|
|
# idx 0: ret=0.01, vol=nan (if min_periods=2) but min_periods=1 => std of [0.01] = nan
|
|
# Actually pandas std() with 1 element is NaN.
|
|
# We should use min_periods=2 or just check the output.
|
|
|
|
df_out = create_labels_volatility(sample_data, returns_window=1, vol_window=3)
|
|
|
|
assert "volatility_label" in df_out.columns
|
|
# Will drop initial NaNs. Std requires at least 2 points.
|
|
# Future returns has 6 rows. The first row's std is NaN. So 5 rows remaining.
|
|
assert len(df_out) == 5
|
|
|
|
# We just ensure it runs and outputs valid values
|
|
assert df_out["volatility_label"].isin([-1, 0, 1]).all()
|
|
assert df_out["volatility"].isna().sum() == 0
|