mirror of
https://github.com/webclinic017/drift.git
synced 2026-08-01 21:27:46 +00:00
30 lines
854 B
Python
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]
|
|
})
|
|
|
|
|
|
|