Files
drift/labeling/labellers/fixed_time_two_class.py
T
2022-01-26 23:22:43 +01:00

30 lines
854 B
Python

from ..types import EventLabeller, EventsDataFrame, ForwardReturnSeries
import pandas as pd
class FixedTimeHorionTwoClassEventLabeller(EventLabeller):
time_horizon: int
def __init__(self, time_horizon: int = 1):
self.time_horizon = time_horizon
def label_events(self, event_start_times: pd.DatetimeIndex, forward_returns: ForwardReturnSeries) -> EventsDataFrame:
event_candidates = forward_returns[event_start_times]
def get_class_binary(x: float) -> int:
return -1 if x <= 0.0 else 1
labels = event_candidates.map(get_class_binary)
return pd.DataFrame({
'start': event_start_times,
'end': event_start_times + pd.Timedelta(days=self.time_horizon),
'label': labels,
'returns': forward_returns[event_start_times]
})