Coverage for src/pyrf/core.py: 100%

1"""Core data structures and utilities for rfpy.""" 

2 

3import attrs 

4import numpy as np 

5import pandas as pd 

6 

7 

8@attrs.define 

9class GoalGrid: 

10 """A grid representing all possible goal combinations between home and away teams. 

11 

12 Creates two DataFrames: 

13 - home: entry [away_goals, home_goals] = home_goals 

14 - away: entry [away_goals, home_goals] = away_goals 

15 

16 Both DataFrames use consistent naming: 

17 - index.name = "away" (represents away team goals) 

18 - columns.name = "home" (represents home team goals) 

19 

20 Attributes: 

21 n: Maximum number of goals for either team (default: 10). 

22 """ 

23 

24 n: int = attrs.field(init=True, repr=True, default=10) 

25 home: pd.DataFrame = attrs.field(repr=False, init=False) 

26 away: pd.DataFrame = attrs.field(repr=False, init=False) 

27 

28 @classmethod 

29 def default(cls) -> GoalGrid: 

30 """Create a default GoalGrid with n=10.""" 

31 return cls(n=10) 

32 

33 def __attrs_post_init__(self) -> None: 

34 """Initialize the home and away goal matrices after object creation.""" 

35 # Assemble home and away goal matrices: 

36 # entry [i,j] = j (home goals) 

37 # entry [i,j] = i (away goals) 

38 nn = np.arange(self.n + 1) 

39 cols = [str(n) for n in nn] 

40 data = np.tile(nn, (self.n + 1, 1)) 

41 self.home = pd.DataFrame(data, index=cols, columns=cols) 

42 self.away = pd.DataFrame(data.T, index=cols, columns=cols) 

43 

44 # DataFrame index is always named "away", columns "home" 

45 self.home.index.name = "away" 

46 self.away.index.name = "away" 

47 self.home.columns.name = "home" 

48 self.away.columns.name = "home"