90 lines
5.0 KiB
Python
90 lines
5.0 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from primethai.indicators import breadth_series, mcclellan, price_features, relative_strength
|
|
from primethai.setups import confirmed_pivots
|
|
|
|
|
|
def bars(n=280):
|
|
idx = pd.bdate_range("2024-01-01", periods=n)
|
|
close = pd.Series(np.linspace(10, 20, n), index=idx)
|
|
return pd.DataFrame({"open": close - .05, "high": close + .25,
|
|
"low": close - .25, "close": close, "volume": 1000.0}, index=idx)
|
|
|
|
|
|
class IndicatorTests(unittest.TestCase):
|
|
def test_ema21_structure_uses_high_close_low_series(self):
|
|
frame = bars()
|
|
result = price_features(frame)
|
|
expected = frame["high"].ewm(span=21, adjust=False, min_periods=21).mean()
|
|
pd.testing.assert_series_equal(result["ema21_high"], expected, check_names=False)
|
|
self.assertEqual(result.iloc[-1]["trend_state"], "UP")
|
|
self.assertEqual(result.iloc[-1]["price_state"], "ABOVE_STRUCTURE")
|
|
|
|
def test_atr_uses_true_range_including_previous_close_gap(self):
|
|
frame = pd.DataFrame({"open": [10, 12, 12], "high": [11, 13, 13],
|
|
"low": [9, 11, 11], "close": [10, 12, 12], "volume": [1, 1, 1]},
|
|
index=pd.bdate_range("2024-01-01", periods=3))
|
|
result = price_features(frame, atr_period=2, adr_period=2)
|
|
self.assertEqual(result["tr"].tolist(), [2, 3, 2])
|
|
self.assertEqual(result["atr14"].iloc[-1], 2.5)
|
|
|
|
def test_breadth_mco_and_mcsi_formulas(self):
|
|
idx = pd.bdate_range("2024-01-01", periods=40)
|
|
closes = pd.DataFrame({"A": np.arange(10, 50), "B": np.arange(50, 10, -1), "C": 10.0}, index=idx)
|
|
membership = {pd.Timestamp(date): ["A", "B", "C"] for date in idx}
|
|
breadth = breadth_series(closes, membership)
|
|
self.assertEqual(breadth.iloc[1]["advancing"], 1)
|
|
self.assertEqual(breadth.iloc[1]["declining"], 1)
|
|
self.assertEqual(breadth.iloc[1]["unchanged"], 1)
|
|
self.assertEqual(breadth.iloc[1]["net_advances"], 0)
|
|
out = mcclellan(breadth, short_span=3, long_span=5, z_window=20)
|
|
net = breadth["net_advances"].astype(float)
|
|
expected = net.ewm(span=3, adjust=False, min_periods=3).mean() - net.ewm(span=5, adjust=False, min_periods=5).mean()
|
|
pd.testing.assert_series_equal(out["mco"], expected, check_names=False)
|
|
self.assertAlmostEqual(out["mcsi"].iloc[-1], out["mco"].fillna(0).cumsum().iloc[-1])
|
|
self.assertEqual(out["mcsi_10dma"].notna().sum(), 31)
|
|
self.assertTrue(set(out["mcsi_state"].dropna()).issubset({"UP", "DOWN", "RECLAIM_10DMA", "LOSE_10DMA", "CURL_UP", "CURL_DOWN"}))
|
|
|
|
def test_rs_percentile_and_set_comparison_are_cross_sectional(self):
|
|
idx = pd.bdate_range("2023-01-02", periods=300)
|
|
benchmark = pd.Series(np.full(300, 100.0), index=idx)
|
|
prices = {"LEADER": pd.Series(np.linspace(10, 40, 300), index=idx),
|
|
"LAGGARD": pd.Series(np.linspace(40, 10, 300), index=idx),
|
|
"FLAT": pd.Series(np.full(300, 20.0), index=idx)}
|
|
scores = relative_strength(prices, benchmark, set_benchmark=benchmark)
|
|
self.assertGreater(scores.loc[idx[-1], "LEADER_rs_percentile"], scores.loc[idx[-1], "FLAT_rs_percentile"])
|
|
self.assertGreater(scores.loc[idx[-1], "FLAT_rs_percentile"], scores.loc[idx[-1], "LAGGARD_rs_percentile"])
|
|
self.assertGreaterEqual(scores.loc[idx[-1], "LEADER_rs_percentile"], 0)
|
|
self.assertLessEqual(scores.loc[idx[-1], "LEADER_rs_percentile"], 100)
|
|
self.assertIn("LEADER_rs_1m_vs_set", scores)
|
|
|
|
def test_rs_percentile_uses_active_point_in_time_members(self):
|
|
idx = pd.bdate_range("2023-01-02", periods=300)
|
|
benchmark = pd.Series(np.full(300, 100.0), index=idx)
|
|
prices = {
|
|
"MEMBER_A": pd.Series(np.linspace(10, 20, 300), index=idx),
|
|
"NONMEMBER": pd.Series(np.linspace(10, 100, 300), index=idx),
|
|
"MEMBER_C": pd.Series(np.full(300, 20.0), index=idx),
|
|
}
|
|
members = {pd.Timestamp(date): ["MEMBER_A", "MEMBER_C"] for date in idx}
|
|
|
|
scores = relative_strength(prices, benchmark, members_by_date=members)
|
|
|
|
# The very strong non-member must not push MEMBER_C below the 50th percentile.
|
|
self.assertAlmostEqual(scores.loc[idx[-1], "MEMBER_C_rs_percentile"], 53.75)
|
|
|
|
def test_higher_low_is_not_published_until_confirmation(self):
|
|
idx = pd.bdate_range("2024-01-01", periods=14)
|
|
lows = pd.Series([10, 9, 8, 9, 10, 11, 10, 9, 10, 11, 12, 13, 14, 15], index=idx)
|
|
pivots = confirmed_pivots(lows, left_bars=2, right_bars=2)
|
|
self.assertTrue(pd.isna(pivots.loc[idx[2], "confirmed_pivot_low"]))
|
|
self.assertTrue(pd.isna(pivots.loc[idx[3], "confirmed_pivot_low"]))
|
|
self.assertEqual(pivots.loc[idx[4], "confirmed_pivot_low"], 8)
|
|
self.assertFalse(bool(pivots.loc[idx[4], "higher_low"]))
|
|
self.assertTrue(pd.isna(pivots.loc[idx[8], "confirmed_pivot_low"]))
|
|
self.assertEqual(pivots.loc[idx[9], "confirmed_pivot_low"], 9)
|
|
self.assertTrue(bool(pivots.loc[idx[9], "higher_low"]))
|