Coverage for gpaw/response/symmetrize.py: 97%

1import numpy as np 

2 

3from gpaw.cgpaw import GG_shuffle 

4 

5from gpaw.response.symmetry import QSymmetries 

6from gpaw.response.qpd import SingleQPWDescriptor 

7 

8 

9class HeadSymmetryOperators: 

10 def __init__(self, symmetries, gd): 

11 self.M_svv = initialize_v_maps(symmetries, gd.cell_cv, gd.icell_cv) 

12 self.sign_s = symmetries.sign_s 

13 self.nsym = len(symmetries) 

14 

15 def symmetrize_wvv(self, A_wvv): 

16 """Symmetrize chi0_wvv""" 

17 tmp_wvv = np.zeros_like(A_wvv) 

18 for M_vv, sign in zip(self.M_svv, self.sign_s): 

19 tmp = np.dot(np.dot(M_vv.T, A_wvv), M_vv) 

20 if sign == 1: 

21 tmp_wvv += np.transpose(tmp, (1, 0, 2)) 

22 elif sign == -1: # transpose head 

23 tmp_wvv += np.transpose(tmp, (1, 2, 0)) 

24 # Overwrite the input 

25 A_wvv[:] = tmp_wvv / self.nsym 

26 

27 

28class BodySymmetryOperators: 

29 def __init__(self, symmetries, qpd): 

30 self.G_sG = initialize_G_maps(symmetries, qpd) 

31 self.sign_s = symmetries.sign_s 

32 self.nsym = len(symmetries) 

33 

34 def symmetrize_wGG(self, A_wGG): 

35 """Symmetrize an array in GG'.""" 

36 for A_GG in A_wGG: 

37 tmp_GG = np.zeros_like(A_GG, order='C') 

38 for G_G, sign in zip(self.G_sG, self.sign_s): 

39 # Numpy: 

40 # if sign == 1: 

41 # tmp_GG += A_GG[G_G, :][:, G_G] 

42 # if sign == -1: 

43 # tmp_GG += A_GG[G_G, :][:, G_G].T 

44 # C: 

45 GG_shuffle(G_G, sign, A_GG, tmp_GG) 

46 A_GG[:] = tmp_GG / self.nsym 

47 

48 # Set up complex frequency alias 

49 symmetrize_zGG = symmetrize_wGG 

50 

51 

52class WingSymmetryOperators(HeadSymmetryOperators): 

53 def __init__(self, symmetries, qpd): 

54 super().__init__(symmetries, qpd.gd) 

55 self.G_sG = initialize_G_maps(symmetries, qpd) 

56 

57 def symmetrize_wxvG(self, A_wxvG): 

58 """Symmetrize chi0_wxvG""" 

59 tmp_wxvG = np.zeros_like(A_wxvG) 

60 for M_vv, sign, G_G in zip(self.M_svv, self.sign_s, self.G_sG): 

61 if sign == 1: 

62 tmp = sign * np.dot(M_vv.T, A_wxvG[..., G_G]) 

63 elif sign == -1: # transpose wings 

64 tmp = sign * np.dot(M_vv.T, A_wxvG[:, ::-1, :, G_G]) 

65 tmp_wxvG += np.transpose(tmp, (1, 2, 0, 3)) 

66 # Overwrite the input 

67 A_wxvG[:] = tmp_wxvG / self.nsym 

68 

69 

70def initialize_v_maps(symmetries: QSymmetries, 

71 cell_cv: np.ndarray, 

72 icell_cv: np.ndarray): 

73 """Calculate cartesian component mapping.""" 

74 return np.array([cell_cv.T @ U_cc.T @ icell_cv 

75 for U_cc in symmetries.U_scc]) 

76 

77 

78def initialize_G_maps(symmetries: QSymmetries, qpd: SingleQPWDescriptor): 

79 """Calculate the Gvector mappings.""" 

80 assert np.allclose(symmetries.q_c, qpd.q_c) 

81 B_cv = 2.0 * np.pi * qpd.gd.icell_cv 

82 G_Gv = qpd.get_reciprocal_vectors(add_q=False) 

83 G_Gc = np.dot(G_Gv, np.linalg.inv(B_cv)) 

84 Q_G = qpd.Q_qG[0] 

85 

86 G_sG = [] 

87 for U_cc, sign, shift_c in symmetries: 

88 iU_cc = np.linalg.inv(U_cc).T 

89 UG_Gc = np.dot(G_Gc - shift_c, sign * iU_cc) 

90 

91 assert np.allclose(UG_Gc.round(), UG_Gc) 

92 UQ_G = np.ravel_multi_index(UG_Gc.round().astype(int).T, 

93 qpd.gd.N_c, 'wrap') 

94 

95 G_G = len(Q_G) * [None] 

96 for G, UQ in enumerate(UQ_G): 

97 try: 

98 G_G[G] = np.argwhere(Q_G == UQ)[0][0] 

99 except IndexError as err: 

100 raise RuntimeError( 

101 'Something went wrong: a symmetry operation mapped a ' 

102 'G-vector outside the plane-wave cutoff sphere') from err 

103 G_sG.append(np.array(G_G, dtype=np.int32)) 

104 return np.array(G_sG)