Coverage for psyop/util.py: 84.44%

1import numpy as np 

2import pandas as pd 

3from rich.table import Table 

4 

5def get_rng(seed: int | np.random.Generator | None) -> np.random.Generator: 

6 """ Get a numpy random number generator from a seed. """ 

7 if isinstance(seed, np.random.Generator): 

8 rng = seed 

9 elif isinstance(seed, int): 

10 rng = np.random.default_rng(seed) 

11 elif seed is None: 

12 rng = np.random.default_rng() 

13 else: 

14 raise TypeError( 

15 f"seed must be None, int, or numpy.random.Generator, not {type(seed)}" 

16 ) 

17 

18 return rng 

19 

20 

21def df_to_table( 

22 pandas_dataframe: pd.DataFrame, 

23 rich_table: Table | None = None, 

24 show_index: bool = True, 

25 index_name: str | None = None, 

26 transpose: bool = True, 

27 sig_figs: int = 3, 

28 heading_style: str = "magenta", 

29) -> Table: 

30 """Convert a pandas.DataFrame to a rich.Table with optional transpose and sig-fig formatting.""" 

31 

32 def _fmt_cell(x) -> str: 

33 try: 

34 import numpy as _np 

35 if isinstance(x, (float, _np.floating)): 

36 if _np.isnan(x) or _np.isinf(x): 

37 return str(x) 

38 return f"{float(x):.{sig_figs}g}" 

39 if isinstance(x, (int, _np.integer)) and not isinstance(x, bool): 

40 return str(int(x)) 

41 except Exception: 

42 pass 

43 return str(x) 

44 

45 rich_table = rich_table or Table(show_header=not transpose, header_style="bold magenta") 

46 

47 df = pandas_dataframe 

48 

49 if not transpose: 

50 # Original orientation 

51 if show_index: 

52 rich_table.add_column(str(index_name) if index_name else "", style=heading_style) 

53 for col in df.columns: 

54 rich_table.add_column(str(col)) 

55 for idx, row in df.iterrows(): 

56 cells = ([str(idx)] if show_index else []) + [_fmt_cell(v) for v in row.tolist()] 

57 rich_table.add_row(*cells) 

58 return rich_table 

59 

60 # Transposed-like view (columns as rows) 

61 left_header = str(index_name) if index_name is not None else "" 

62 rich_table.add_column(left_header, style=heading_style) # magenta left column 

63 

64 # Column headers across (not displayed when show_header=False, but keeps widths aligned) 

65 across_headers = [str(i) for i in df.index] if show_index else ["" for _ in range(len(df.index))] 

66 for h in across_headers: 

67 rich_table.add_column(h) 

68 

69 for col in df.columns: 

70 row_vals = [_fmt_cell(v) for v in df[col].tolist()] 

71 rich_table.add_row(str(col), *row_vals) 

72 

73 return rich_table