Coverage for hiphive/cluster_space_data.py: 100%

1from typing import Any, Dict, List 

2import numpy as np 

3from .input_output.pretty_table_prints import print_table 

4from pandas import DataFrame 

5 

6 

7class ClusterSpaceData: 

8 

9 """Container class that holds information concerning a cluster space, including, 

10 e.g., the number of cluster, orbits, and parameters by order and number of bodies. 

11 

12 Parameters 

13 ---------- 

14 cs : ClusterSpace 

15 cluster space 

16 """ 

17 

18 def __init__(self, cs) -> None: 

19 self.max_nbody = cs.cutoffs.max_nbody 

20 self.max_order = cs.cutoffs.max_order 

21 

22 # collect cutoff matrix 

23 self._cutoff_matrix = cs.cutoffs.cutoff_matrix 

24 

25 # collect cluster, orbit, eigentensor counts 

26 self._cluster_counts = np.zeros((self.max_nbody, self.max_order - 1), dtype=int) 

27 self._orbit_counts = np.zeros((self.max_nbody, self.max_order - 1), dtype=int) 

28 self._eigentensor_counts = np.zeros((self.max_nbody, self.max_order - 1), dtype=int) 

29 for orbit in cs.orbits: 

30 proto_cluster = cs.cluster_list[orbit.prototype_index] 

31 order = len(proto_cluster) 

32 nbody = len(set(proto_cluster)) 

33 self._cluster_counts[nbody-1, order-2] += len(orbit.orientation_families) 

34 self._orbit_counts[nbody-1, order-2] += 1 

35 self._eigentensor_counts[nbody-1, order-2] += len(orbit.eigentensors) 

36 

37 # collect number of parameters after sum rule constraint 

38 self.ndofs_by_order = {o: cs.get_n_dofs_by_order(o) for o in cs.cutoffs.orders} 

39 

40 def print_tables(self) -> None: 

41 """ Prints information concerning the underlying cluster space to stdout, including, 

42 e.g., the number of cluster, orbits, and parameters by order and number of bodies. """ 

43 

44 # print table data 

45 print('Cutoff Matrix') 

46 print_table(self._cutoff_matrix_padded) 

47 print('\nCluster counts') 

48 print_table(self._cluster_counts, include_sum=True) 

49 print('\nOrbit counts') 

50 print_table(self._orbit_counts, include_sum=True) 

51 print('\nEigentensor counts') 

52 print_table(self._eigentensor_counts, include_sum=True) 

53 

54 @property 

55 def _cutoff_matrix_padded(self) -> np.array: 

56 """ Padded cutoff matrix with None for nbody=1 terms """ 

57 return np.vstack(([[None] * (self.max_order - 1)], self._cutoff_matrix)) 

58 

59 def to_list(self) -> List[Dict[str, Any]]: 

60 """ Returns cluster space data in the form of a list of dicts. """ 

61 records = [] 

62 for order in range(2, self.max_order+1): 

63 for nbody in range(1, self.max_nbody+1): 

64 if nbody > order: 

65 continue 

66 row = dict(order=order, nbody=nbody) 

67 row['cutoff'] = self._cutoff_matrix_padded[nbody-1, order-2] 

68 row['cluster_counts'] = self._cluster_counts[nbody-1, order-2] 

69 row['orbit_counts'] = self._orbit_counts[nbody-1, order-2] 

70 row['eigentensor_counts'] = self._eigentensor_counts[nbody-1, order-2] 

71 records.append(row) 

72 return records 

73 

74 def to_dataframe(self) -> DataFrame: 

75 """ Returns cluster space data in the form of a pandas DataFrame. """ 

76 return DataFrame.from_dict(self.to_list()) 

77 

78 def __str__(self) -> str: 

79 s = f'Cutoff matrix: \n {self._cutoff_matrix}\n' 

80 s += f'Cluster counts: \n{self._cluster_counts}\n' 

81 s += f'Orbit counts: \n{self._orbit_counts}\n' 

82 s += f'Eigentensor counts: \n{self._eigentensor_counts}\n' 

83 s += f'Degrees of freedoms: \n{self.ndofs_by_order}\n' 

84 return s 

85 

86 def __repr__(self) -> str: 

87 return str(self)