Coverage for adrs/alpha.py: 29%

1import polars as pl 

2from abc import abstractmethod 

3from datetime import datetime, timedelta 

4from typing import Any, TypedDict, Unpack, NotRequired, cast 

5 

6from adrs.types import Performance 

7from adrs.performance import Evaluator 

8from adrs.data import Datamap, DataInfo, DataProcessor 

9from adrs.performance.metric import Ratio, Drawdown, Trade 

10 

11 

12class AlphaBacktestArgs(TypedDict): 

13 evaluator: Evaluator 

14 base_asset: str 

15 datamap: Datamap 

16 start_time: datetime 

17 end_time: datetime 

18 fees: float 

19 data_df: NotRequired[pl.DataFrame] 

20 price_shift: NotRequired[int] 

21 output_columns: NotRequired[list[pl.Expr]] 

22 

23 

24class Alpha: 

25 def __init__( 

26 self, 

27 id: str, 

28 data_infos: list[DataInfo], 

29 data_processor: DataProcessor, 

30 ): 

31 self.id = id 

32 self.data_infos = data_infos 

33 self.data_processor = data_processor 

34 self.data_processor.data_infos = data_infos 

35 

36 @abstractmethod 

37 def next(self, data_df: pl.DataFrame) -> pl.DataFrame: 

38 raise NotImplementedError("Every alpha must implement its `next` method") 

39 

40 def backtest( 

41 self, **kwargs: Unpack[AlphaBacktestArgs] 

42 ) -> tuple[Performance, pl.DataFrame]: 

43 evaluator = kwargs["evaluator"] 

44 base_asset = kwargs["base_asset"] 

45 datamap = kwargs["datamap"] 

46 start_time = kwargs["start_time"] 

47 end_time = kwargs["end_time"] 

48 fees = kwargs["fees"] 

49 data_df = kwargs.get("data_df", self.data_processor.process(datamap)) 

50 price_shift = kwargs.get("price_shift", 0) 

51 output_columns = kwargs.get("output_columns", [pl.all()]) 

52 

53 if data_df is None: 

54 raise ValueError("data_df received is None") 

55 

56 df = self.next(data_df) 

57 interval = cast(timedelta, df["start_time"].diff().last()) 

58 

59 # Verify that the signal is valid 

60 if "signal" not in df.schema.keys(): 

61 raise ValueError( 

62 "`next()` method must return a DataFrame with the column 'signal'" 

63 ) 

64 if not df.schema["signal"].is_numeric(): 

65 raise ValueError( 

66 "DataFrame returned from `next()` must have a 'signal' column that is numeric" 

67 ) 

68 if ( 

69 df.select("signal").min().to_numpy().ravel()[0] < -1 

70 or df.select("signal").max().to_numpy().ravel()[0] > 1 

71 ): 

72 raise ValueError( 

73 "DataFrame returned from `next()` must have a 'signal' column that is between [-1, 1]" 

74 ) 

75 

76 pdf = evaluator.eval( 

77 signal_lf=df.lazy(), 

78 base_asset=base_asset, 

79 datamap=datamap, 

80 start_time=start_time, 

81 end_time=end_time, 

82 fees=fees, 

83 interval=interval, 

84 price_shift=price_shift, 

85 output_columns=output_columns, 

86 ).collect(engine="in-memory") 

87 

88 # Compute the metrics 

89 performance: dict[str, Any] = { 

90 "start_time": start_time, 

91 "end_time": end_time, 

92 "resolution_detected": interval, 

93 "offset_detected": evaluator.last_grid_phase, 

94 "metadata": {}, 

95 } 

96 for metric in [Ratio(interval=interval), Trade(), Drawdown(interval=interval)]: 

97 result = metric.compute(pdf) 

98 performance = {**performance, **result} 

99 

100 return Performance.model_validate(performance), pdf