Coverage for adrs/performance/evaluator.py: 91%

1import logging 

2import polars as pl 

3from datetime import datetime, timedelta, timezone 

4 

5from adrs.types import Topic 

6from adrs.data.datamap import DataInfo, Datamap 

7 

8logger = logging.getLogger(__name__) 

9 

10 

11def grid_phase( 

12 first_start_time: datetime | None, interval: str | timedelta 

13) -> timedelta: 

14 """ 

15 Get offset of signal for group_by_dynamic(offset=...) 

16 """ 

17 if first_start_time is None or not isinstance(interval, timedelta): 

18 return timedelta(0) 

19 interval_s = interval.total_seconds() 

20 if interval_s <= 0: 

21 return timedelta(0) 

22 if first_start_time.tzinfo is None: 

23 first_start_time = first_start_time.replace(tzinfo=timezone.utc) 

24 return timedelta(seconds=first_start_time.timestamp() % interval_s) 

25 

26 

27def phase_drift(start_times: pl.Series, interval: str | timedelta) -> list[timedelta]: 

28 """Distinct grid phases in `start_times`, empty when the phase is constant.""" 

29 if not isinstance(interval, timedelta) or start_times.len() < 2: 

30 return [] 

31 interval_us = interval // timedelta(microseconds=1) 

32 if interval_us <= 0: 

33 return [] 

34 phases = start_times.dt.epoch("us") % interval_us 

35 if phases.n_unique() < 2: 

36 return [] 

37 return [timedelta(microseconds=int(p)) for p in phases.unique().sort()] 

38 

39 

40def has_column(info: DataInfo, col: str) -> bool: 

41 for column in info.columns: 

42 if column.dst == col: 

43 return True 

44 return False 

45 

46 

47class Evaluator: 

48 def __init__(self, assets: dict[str, DataInfo]): 

49 for k, v in assets.items(): 

50 if not has_column(info=v, col="price"): 

51 raise ValueError(f"Asset {k} must have a 'dst' column of 'price'") 

52 

53 self.assets = assets 

54 self.last_grid_phase: timedelta | None = None 

55 

56 def eval( 

57 self, 

58 signal_lf: pl.LazyFrame, 

59 base_asset: str, 

60 datamap: Datamap, 

61 start_time: datetime, 

62 end_time: datetime, 

63 fees: float, 

64 interval: str | timedelta, 

65 price_shift: int = 0, 

66 output_columns: list[pl.Expr] = [pl.all()], 

67 ): 

68 if price_shift < 0: 

69 raise ValueError("price_shift must be non-negative") 

70 

71 if base_asset not in self.assets: 

72 raise ValueError(f"Base asset {base_asset} not found in configured assets") 

73 

74 info = self.assets[base_asset] 

75 if Topic.from_str(info.topic) not in datamap.keys(): 

76 raise ValueError(f"Data for base asset {base_asset} not found in datamap") 

77 

78 prices_lf = ( 

79 datamap.get(info).lazy().with_columns(pl.col("price").shift(-price_shift)) 

80 ) 

81 

82 start_times = signal_lf.select(pl.col("start_time")).collect()["start_time"] 

83 phase = grid_phase(start_times.min(), interval) 

84 

85 if drift := phase_drift(start_times, interval): 

86 shown = ", ".join(str(p) for p in drift[:5]) 

87 more = f" (+{len(drift) - 5} more)" if len(drift) > 5 else "" 

88 logger.error( 

89 "Signal grid phase for base asset %s is not constant: %d distinct " 

90 "phases [%s]%s against interval %s. Prices can only be resampled onto " 

91 "one phase (%s), so rows on the others would be dropped by the join " 

92 "and silently zeroed.", 

93 base_asset, 

94 len(drift), 

95 shown, 

96 more, 

97 interval, 

98 phase, 

99 ) 

100 raise ValueError( 

101 f"Signal grid phase for base asset {base_asset} drifts mid-history: " 

102 f"{len(drift)} distinct phases [{shown}]{more} for interval {interval}" 

103 ) 

104 

105 self.last_grid_phase = phase 

106 

107 df = ( 

108 prices_lf.group_by_dynamic( 

109 index_column="start_time", every=interval, offset=phase 

110 ) 

111 .agg(pl.col("price").last()) 

112 .drop_nulls() 

113 .join(signal_lf, how="left", on="start_time") 

114 .filter( 

115 pl.col("start_time").is_between(start_time, end_time, closed="left") 

116 ) 

117 .with_columns( 

118 pl.col("signal").forward_fill().fill_null(strategy="zero"), 

119 pl.col("signal") 

120 .shift(1) 

121 .alias("prev_signal") 

122 .forward_fill() 

123 .fill_null(strategy="zero"), 

124 pl.col("price") 

125 .pct_change() 

126 .alias("returns") 

127 .fill_null(strategy="zero"), 

128 ) 

129 .with_columns( 

130 (pl.col("signal") - pl.col("prev_signal")) 

131 .alias("trade") 

132 .fill_null(strategy="zero") 

133 ) 

134 .with_columns( 

135 ( 

136 pl.col("prev_signal") * pl.col("returns") 

137 - pl.col("trade").abs() * fees / 100 

138 ) 

139 .alias("pnl") 

140 .fill_null(strategy="zero"), 

141 ) 

142 .with_columns( 

143 pl.col("pnl").cum_sum().alias("equity").fill_null(strategy="zero") 

144 ) 

145 ) 

146 return df.select(*output_columns)