Coverage for src/duelboard/calculators/base.py: 96%

1"""Base Elo rating calculator using online linear update algorithm.""" 

2 

3from collections import defaultdict 

4 

5import pandas as pd 

6 

7from duelboard.models import Battle, BattleOutcome, EloRating 

8from duelboard.types import RatingsDict 

9 

10 

11class EloCalculator: 

12 """Base Elo rating calculator using online linear update algorithm.""" 

13 

14 def __init__( 

15 self, 

16 k_factor: float = 4, 

17 scale: float = 400, 

18 base: float = 10, 

19 initial_rating: float = 1000, 

20 ) -> None: 

21 """Initialize the Elo calculator. 

22 

23 Args: 

24 k_factor: K-factor for Elo updates (lower = more stable) 

25 scale: Scale parameter for Elo calculation 

26 base: Base for exponential calculation 

27 initial_rating: Initial rating for new players 

28 """ 

29 self.k_factor = k_factor 

30 self.scale = scale 

31 self.base = base 

32 self.initial_rating = initial_rating 

33 

34 def calculate(self, battles: list[Battle] | pd.DataFrame) -> RatingsDict: 

35 """Calculate Elo ratings from a list of battles. 

36 

37 Args: 

38 battles: List of Battle objects or DataFrame with battles 

39 

40 Returns: 

41 Dictionary mapping player names to EloRating objects 

42 """ 

43 if isinstance(battles, pd.DataFrame): 

44 battles = self._dataframe_to_battles(battles) 

45 

46 ratings = defaultdict(lambda: self.initial_rating) 

47 battle_counts = defaultdict(int) 

48 

49 for battle in battles: 

50 player_a = battle.player_a 

51 player_b = battle.player_b 

52 outcome = battle.outcome 

53 

54 rating_a = ratings[player_a] 

55 rating_b = ratings[player_b] 

56 

57 expected_a = self._calculate_expected_score(rating_a, rating_b) 

58 expected_b = 1 - expected_a 

59 

60 actual_a = self._outcome_to_score(outcome, is_player_a=True) 

61 actual_b = 1 - actual_a 

62 

63 new_rating_a = rating_a + self.k_factor * (actual_a - expected_a) 

64 new_rating_b = rating_b + self.k_factor * (actual_b - expected_b) 

65 

66 ratings[player_a] = new_rating_a 

67 ratings[player_b] = new_rating_b 

68 

69 battle_counts[player_a] += 1 

70 battle_counts[player_b] += 1 

71 

72 return { 

73 player: EloRating( 

74 player=player, 

75 rating=rating, 

76 battles=battle_counts[player], 

77 ) 

78 for player, rating in ratings.items() 

79 } 

80 

81 def _calculate_expected_score(self, rating_a: float, rating_b: float) -> float: 

82 """Calculate expected score for player A.""" 

83 return 1 / (1 + self.base ** ((rating_b - rating_a) / self.scale)) 

84 

85 def _outcome_to_score(self, outcome: BattleOutcome, *, is_player_a: bool) -> float: 

86 """Convert battle outcome to score for player A or B.""" 

87 if outcome in (BattleOutcome.TIE, BattleOutcome.TIE_BOTHBAD): 

88 return 0.5 

89 

90 if is_player_a: 

91 return 1.0 if outcome == BattleOutcome.WIN_A else 0.0 

92 return 1.0 if outcome == BattleOutcome.WIN_B else 0.0 

93 

94 def _dataframe_to_battles(self, df: pd.DataFrame) -> list[Battle]: 

95 """Convert DataFrame to list of Battle objects.""" 

96 battles = [] 

97 for _, row in df.iterrows(): 

98 # Try different column name formats for compatibility 

99 player_a = row.get("player_a", row.get("player_a")) 

100 player_b = row.get("player_b", row.get("player_b")) 

101 winner = row.get("winner", row.get("result")) 

102 

103 battle = Battle( 

104 player_a=player_a, 

105 player_b=player_b, 

106 outcome=BattleOutcome(winner), 

107 metadata=row.to_dict() if hasattr(row, "to_dict") else {}, 

108 ) 

109 battles.append(battle) 

110 return battles 

111 

112 def predict_win_probability( 

113 self, 

114 player_a: str, 

115 player_b: str, 

116 ratings: RatingsDict, 

117 ) -> float: 

118 """Predict win probability for player A against player B. 

119 

120 Args: 

121 player_a: Name of player A 

122 player_b: Name of player B 

123 ratings: Dictionary of current ratings 

124 

125 Returns: 

126 Probability that player A wins 

127 """ 

128 rating_a = ratings[player_a].rating 

129 rating_b = ratings[player_b].rating 

130 return self._calculate_expected_score(rating_a, rating_b) 

131 

132 def get_leaderboard(self, ratings: RatingsDict) -> list[EloRating]: 

133 """Get sorted leaderboard from ratings. 

134 

135 Args: 

136 ratings: Dictionary of ratings 

137 

138 Returns: 

139 List of EloRating objects sorted by rating (descending) 

140 """ 

141 return sorted(ratings.values(), key=lambda x: x.rating, reverse=True)