You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
|
|
_SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "enrich_catalyst_persistence_features.py"
|
|
_SPEC = importlib.util.spec_from_file_location("enrich_catalyst_persistence_features", _SCRIPT_PATH)
|
|
assert _SPEC and _SPEC.loader
|
|
_MODULE = importlib.util.module_from_spec(_SPEC)
|
|
sys.modules[_SPEC.name] = _MODULE
|
|
_SPEC.loader.exec_module(_MODULE)
|
|
compute_catalyst_persistence_features = _MODULE.compute_catalyst_persistence_features
|
|
|
|
|
|
def test_compute_catalyst_persistence_features_counts_only_prior_same_ticker_events() -> None:
|
|
rows = pd.DataFrame(
|
|
[
|
|
{
|
|
"ticker": "AAA",
|
|
"event_date": "2022-01-10",
|
|
"event_type": "earnings_release",
|
|
"__split_name": "train",
|
|
"__split_order": 0,
|
|
"__split_row_idx": 0,
|
|
"__event_date_obj": pd.Timestamp("2022-01-10").date(),
|
|
},
|
|
{
|
|
"ticker": "AAA",
|
|
"event_date": "2022-01-25",
|
|
"event_type": "management_change",
|
|
"__split_name": "train",
|
|
"__split_order": 0,
|
|
"__split_row_idx": 1,
|
|
"__event_date_obj": pd.Timestamp("2022-01-25").date(),
|
|
},
|
|
{
|
|
"ticker": "AAA",
|
|
"event_date": "2022-02-20",
|
|
"event_type": "material_contract",
|
|
"__split_name": "valid",
|
|
"__split_order": 1,
|
|
"__split_row_idx": 0,
|
|
"__event_date_obj": pd.Timestamp("2022-02-20").date(),
|
|
},
|
|
{
|
|
"ticker": "BBB",
|
|
"event_date": "2022-02-20",
|
|
"event_type": "earnings_release",
|
|
"__split_name": "valid",
|
|
"__split_order": 1,
|
|
"__split_row_idx": 1,
|
|
"__event_date_obj": pd.Timestamp("2022-02-20").date(),
|
|
},
|
|
]
|
|
)
|
|
|
|
enriched = compute_catalyst_persistence_features(rows)
|
|
second = enriched.loc[(enriched["ticker"] == "AAA") & (enriched["event_date"] == "2022-01-25")].iloc[0]
|
|
third = enriched.loc[(enriched["ticker"] == "AAA") & (enriched["event_date"] == "2022-02-20")].iloc[0]
|
|
other = enriched.loc[enriched["ticker"] == "BBB"].iloc[0]
|
|
|
|
assert second["prior_catalyst_count_20d"] == 1.0
|
|
assert second["prior_catalyst_type_diversity_20d"] == 1.0
|
|
assert third["prior_catalyst_count_20d"] == 0.0
|
|
assert third["prior_catalyst_count_60d"] == 2.0
|
|
assert third["prior_catalyst_type_diversity_60d"] == 2.0
|
|
assert other["prior_catalyst_count_60d"] == 0.0
|
|
assert other["prior_catalyst_type_diversity_60d"] == 0.0
|