Coverage for gpaw/atomrotations.py: 100%

1import numpy as np 

2 

3from gpaw.rotation import rotation 

4from gpaw.utilities import pack_density, unpack_density 

5 

6 

7class SingleAtomRotations: 

8 def __init__(self, R_sii): 

9 self.R_sii = R_sii 

10 

11 @classmethod 

12 def new(cls, ni, l_j, R_slmm): 

13 nsym = len(R_slmm) 

14 R_sii = np.zeros((nsym, ni, ni)) 

15 i1 = 0 

16 for l in l_j: 

17 i2 = i1 + 2 * l + 1 

18 for s, R_lmm in enumerate(R_slmm): 

19 R_sii[s, i1:i2, i1:i2] = R_lmm[l] 

20 i1 = i2 

21 return cls(R_sii) 

22 

23 def symmetrize(self, a, D_aii, map_sa): 

24 ni = self.R_sii.shape[1] 

25 D_ii = np.zeros((ni, ni)) 

26 for s, R_ii in enumerate(self.R_sii): 

27 D_ii += R_ii @ D_aii[map_sa[s][a]] @ R_ii.T 

28 return D_ii / len(map_sa) 

29 

30 

31class AtomRotations: 

32 def __init__(self, setups, id_a, symmetry): 

33 R_slmm = [] 

34 for op_cc in symmetry.op_scc: 

35 op_vv = np.linalg.inv(symmetry.cell_cv) @ op_cc @ symmetry.cell_cv 

36 R_slmm.append([rotation(l, op_vv) for l in range(4)]) 

37 

38 rotations = {} 

39 for key, setup in setups.items(): 

40 rotations[key] = SingleAtomRotations.new(setup.ni, setup.l_j, 

41 R_slmm) 

42 

43 self._rotations = rotations 

44 self._id_a = id_a 

45 

46 def get_R_asii(self): 

47 return [self.get_by_a(a).R_sii for a in range(len(self._id_a))] 

48 

49 def get_by_a(self, a): 

50 return self._rotations[self._id_a[a]] 

51 

52 def symmetrize_atomic_density_matrices(self, D_asp, a_sa): 

53 if not D_asp: 

54 return 

55 

56 nspins = next(iter(D_asp.values())).shape[0] 

57 

58 for s in range(nspins): 

59 D_aii = [unpack_density(D_asp[a][s]) for a in range(len(D_asp))] 

60 for a, D_ii in enumerate(D_aii): 

61 D_asp[a][s] = pack_density( 

62 self.get_by_a(a).symmetrize(a, D_aii, a_sa))