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

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

2 

3from typing import Any 

4 

5import numpy as np 

6import pandas as pd 

7 

8 

9def calc_total_mass(df: pd.DataFrame) -> float: 

10 """Calculate total probability mass.""" 

11 result: Any = df.sum().sum() 

12 return float(result) 

13 

14 

15def sum_diagonals(npy: np.ndarray) -> pd.Series: 

16 """Sum the diagonals of a 2D numpy array representing a PMF, yielding goal differences.""" 

17 m, n = npy.shape 

18 data = {k: npy.diagonal(offset=k).sum() for k in range(-m + 1, n)} 

19 return pd.Series(data, dtype=float)