Coverage for hiphive/core/translational_constraints.py: 99%

94 statements  

« prev     ^ index     » next       coverage.py v7.6.8, created at 2024-11-28 11:20 +0000

1""" 

2Functionality for enforcing translational sum rules 

3""" 

4 

5import sympy 

6from scipy.sparse import coo_matrix 

7import numpy as np 

8 

9from ..input_output.logging_tools import logger 

10from .utilities import SparseMatrix, BiMap 

11from .orbits import get_orbits 

12 

13 

14logger = logger.getChild('sum_rules') 

15 

16 

17def create_constraint_map(cs): 

18 

19 # First we need to know the global index of a eigentensor given which 

20 # orbit it belongs to. 

21 params = [] 

22 n = 0 

23 for orbit_index, orbit in enumerate(cs.orbits): 

24 nParams_in_orbit = len(orbit.eigentensors) 

25 params.append(list(range(n, n + nParams_in_orbit))) 

26 n += nParams_in_orbit 

27 

28 # If we dont want the sum rules an eye matrix is created instead. 

29 if not cs.acoustic_sum_rules: 

30 row = list(range(n)) 

31 col = list(range(n)) 

32 data = [1.0]*n 

33 cs._cvs = coo_matrix((data, (row, col)), shape=(n, n)) 

34 return None 

35 

36 M = get_translational_constraint_matrix(cs) 

37 

38 nullspace = M.nullspace() 

39 logger.debug('Assembling solutions...') 

40 row = [] 

41 col = [] 

42 data = [] 

43 for c, vec in enumerate(nullspace): 

44 for r, _, v in vec.row_list(): 

45 row.append(r) 

46 col.append(c) 

47 data.append(np.float64(v)) 

48 

49 # sanity check that there are orbits still present 

50 if len(nullspace) == 0: 

51 raise ValueError('There are no degrees of freedom, error!') 

52 

53 # This is the final product, the constraint map (or constraint vectors) 

54 cs._cvs = coo_matrix((data, (row, col)), shape=(len(vec), len(nullspace))) 

55 

56 logger.debug('{} degrees of freedom'.format(cs._cvs.shape[1])) 

57 logger.debug('Finished constructing constraint vectors') 

58 

59 

60def get_translational_constraint_matrix(cs): 

61 logger.debug('Starting constructing constraint matrices') 

62 

63 # First we need to know the global index of a eigentensor given which 

64 # orbit it belongs to. 

65 params = [] 

66 n = 0 

67 for orbit_index, orbit in enumerate(cs.orbits): 

68 nParams_in_orbit = len(orbit.eigentensors) 

69 params.append(list(range(n, n + nParams_in_orbit))) 

70 n += nParams_in_orbit 

71 

72 # The lookup is a fast way to get the eigentensors given a cluster Also the 

73 # orbit index is bundled to quickly find the correct parameter indices 

74 lookup = {} 

75 for orbit_index, orbit in enumerate(cs.orbits): 

76 for of in orbit.orientation_families: 

77 for cluster_index, perm_index in zip(of.cluster_indices, 

78 of.permutation_indices): 

79 cluster = cs._cluster_list[cluster_index] 

80 perm = cs._permutations[perm_index] 

81 lookup[tuple(cluster)] = [et.transpose(perm) for et in 

82 of.eigentensors], orbit_index 

83 

84 # Init just one row to be able to stack. n is the total number of 

85 # eigentensors 

86 M = SparseMatrix(1, n, 0) # This is the constraint matrix 

87 for order in cs.cutoffs.orders: 

88 logger.debug('Order {}: '.format(order)) 

89 

90 # Find which minimal amount of prefixes that is needed 

91 # Essentially they are all the prototypes of order one less still 

92 # compatible with the cutoff. 

93 prefix_list = [] 

94 if order == 2: 

95 for i in np.unique(cs.wyckoff_sites): 

96 prefix_list.append((i,)) 

97 else: 

98 tmp_cluster_set = set() 

99 for cluster in cs._cluster_list: 

100 if len(cluster) != order: 

101 continue 

102 for i in range(order): 

103 c = tuple(cluster[:i]) + tuple(cluster[i+1:]) 

104 if c[0] < len(cs._prim): 

105 tmp_cluster_set.add(c) 

106 tmp_cluster_list = BiMap() 

107 for c in tmp_cluster_set: 

108 tmp_cluster_list.append(c) 

109 orbits = get_orbits(tmp_cluster_list, cs._atom_list, 

110 cs.rotation_matrices, cs.translation_vectors, 

111 cs.permutations, cs._prim, cs.symprec) 

112 for orb in orbits: 

113 prefix_list.append(tmp_cluster_list[orb.prototype_index]) 

114 

115 m = SparseMatrix(3**order, M.shape[1], 0) # intermediate matrix 

116 # Now loop over the prefix and add the last atom index 

117 for prefix in prefix_list: 

118 m *= 0 

119 cluster = list(prefix) + [None] 

120 for i in range(len(cs._atom_list)): 

121 cluster[-1] = i 

122 tmp = lookup.get(tuple(sorted(cluster))) 

123 if tmp is None: 

124 continue 

125 ets, orbit_index = tmp 

126 inv_argsort = np.argsort(np.argsort(cluster)) 

127 ets = [et.transpose(inv_argsort) for et in ets] 

128 for et, col in zip(ets, params[orbit_index]): 

129 et_flat = et.flatten() 

130 for j in et_flat.nonzero()[0]: 

131 m[j, col] += sympy.nsimplify(et_flat[j]) 

132 M = SparseMatrix.vstack(M, m) 

133 # only compress if there are very many elements 

134 if M.nnz() > cs._config['max_number_constraint_elements']: 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true

135 M = SparseMatrix.rref(M)[0] 

136 

137 logger.debug('Finished constructing constraint matrices') 

138 logger.debug('Starting constructing constraint vectors') 

139 return M