Coverage for gpaw/analyse/wignerseitz.py: 61%

1from math import sqrt, pi 

2 

3import numpy as np 

4from ase.units import Bohr 

5 

6from gpaw.utilities import pack_density 

7from gpaw.analyse.hirshfeld import HirshfeldDensity 

8from gpaw.utilities.tools import coordinates 

9 

10 

11def wignerseitz(gd, atoms, scale=None): 

12 """Determine which atom is closest to each grid point. 

13 

14 The atom distances might be scaled by the scale list.""" 

15 if scale is None: 

16 scale = [1.] * len(atoms) 

17 else: 

18 assert len(scale) == len(atoms) 

19 r_vG, R2min_G = coordinates(gd, atoms[0].position / Bohr) 

20 R2min_G *= scale[0] ** 2 

21 index_G = gd.zeros(dtype=int) 

22 for a, atom in enumerate(atoms[1:]): 

23 r_vG, r2_G = coordinates(gd, atom.position / Bohr) 

24 r2_G *= scale[a + 1] ** 2 

25 index_G = np.where(R2min_G > r2_G, a + 1, index_G) 

26 R2min_G = np.where(R2min_G > r2_G, r2_G, R2min_G) 

27 return index_G 

28 

29 

30class WignerSeitz: 

31 

32 def __init__(self, gd, atoms, 

33 calculator=None, scale=None): 

34 """Find the grid points nearest to the atoms""" 

35 

36 self.atoms = atoms 

37 self.gd = gd 

38 self.calculator = calculator 

39 

40 self.atom_index = wignerseitz(gd, atoms, scale) 

41 

42 def expand(self, density): 

43 """Expand a smooth density in Wigner-Seitz cells around the atoms""" 

44 n = len(self.atoms) 

45 weights = np.empty((n,)) 

46 for a in range(n): 

47 mask = np.where(self.atom_index == a, density, 0.0) 

48 # XXX Optimize! No need to integrate in zero-region 

49 weights[a] = self.gd.integrate(mask) 

50 

51 return weights 

52 

53 def expand_density(self, nt_G, s, nspins): 

54 """Get the weights of spin-density in Wigner-Seitz cells 

55 around the atoms. The spin index and number of spins are 

56 needed for the augmentation sphere corrections.""" 

57 weights_a = self.expand(nt_G) 

58 for a, nucleus in enumerate(self.atoms): 

59 weights_a[a] += nucleus.get_density_correction(s, nspins) 

60 return weights_a 

61 

62 def expand_wave_function(self, psit_G, u, n): 

63 """Get the weights of wave function in Wigner-Seitz cells 

64 around the atoms. The spin-k-point index u and band number n 

65 are needed for the augmentation sphere corrections.""" 

66 

67 assert psit_G.dtype == float 

68 # smooth part 

69 weigths = self.expand(psit_G ** 2) 

70 

71 # add augmentation sphere corrections 

72 for a, nucleus in enumerate(self.atoms): 

73 P_i = nucleus.P_uni[u, n] 

74 P_p = pack_density(np.outer(P_i, P_i)) 

75 Delta_p = sqrt(4 * pi) * nucleus.setup.Delta_pL[:, 0] 

76 weigths[a] += np.dot(Delta_p, P_p) 

77 

78 return weigths 

79 

80 def get_charges(self, den_g): 

81 """Charge on the atom according to the Wigner-Seitz partitioning 

82 

83 Can be applied to any density den_g. 

84 """ 

85 assert den_g.shape == tuple(self.gd.n_c) 

86 charges = [] 

87 for atom, q in zip(self.atoms, self.expand(den_g)): 

88 charges.append(atom.number - q) 

89 return charges 

90 

91 def get_effective_volume_ratio(self, atom_index): 

92 """Effective volume to free volume ratio. 

93 

94 After: Tkatchenko and Scheffler PRL 102 (2009) 073005 

95 """ 

96 atoms = self.atoms 

97 finegd = self.gd 

98 

99 den_g, gd = self.calculator.density.get_all_electron_density(atoms) 

100 assert gd == finegd 

101 denfree_g, gd = self.hdensity.get_density([atom_index]) 

102 assert gd == finegd 

103 

104 # the atoms r^3 grid 

105 position = self.atoms[atom_index].position / Bohr 

106 r_vg, r2_g = coordinates(finegd, origin=position) 

107 r3_g = r2_g * np.sqrt(r2_g) 

108 

109 weight_g = np.where(self.atom_index == atom_index, 1.0, 0.0) 

110 

111 nom = finegd.integrate(r3_g * den_g[0] * weight_g) 

112 denom = finegd.integrate(r3_g * denfree_g) 

113 

114 return nom / denom 

115 

116 def get_effective_volume_ratios(self): 

117 """Return the list of effective volume to free volume ratios.""" 

118 ratios = [] 

119 self.hdensity = HirshfeldDensity(self.calculator) 

120 for a, atom in enumerate(self.atoms): 

121 ratios.append(self.get_effective_volume_ratio(a)) 

122 return np.array(ratios) 

123 

124 

125class LDOSbyBand: 

126 

127 """Base class for a band by band LDOS""" 

128 

129 def by_element(self): 

130 # get element indicees 

131 elemi = {} 

132 for i, nucleus in enumerate(self.paw.atoms): 

133 symbol = nucleus.setup.symbol 

134 if symbol in elemi: 

135 elemi[symbol].append(i) 

136 else: 

137 elemi[symbol] = [i] 

138 for key in elemi.keys(): 

139 elemi[key] = self.get(elemi[key]) 

140 return elemi 

141 

142 

143''' 

144class WignerSeitzLDOS(LDOSbyBand): 

145 

146 """Class to get the unfolded LDOS defined by Wigner-Seitz cells""" 

147 

148 def __init__(self, paw): 

149 self.paw = paw 

150 self.ws = WignerSeitz(paw.gd, paw.atoms) 

151 

152 nu = paw.nkpts * paw.nspins 

153 ldos = np.empty((nu, paw.nbands, len(paw.atoms))) 

154 for u, kpt in enumerate(paw.kpt_u): 

155 for n, psit_G in enumerate(kpt.psit_nG): 

156 # XXX variable ws undefined. Probably self.ws 

157 ldos[u, n, :] = ws.expand_wave_function(psit_G, u, n) 

158 

159 def write(self, filename, slavewrite=False): 

160 if self.world.rank == MASTER or slavewrite: 

161 paw = self.paw 

162 f = open(filename, 'w') 

163 

164 nn = len(paw.atoms) 

165 

166 for k in range(paw.nkpts): 

167 for s in range(paw.nspins): 

168 u = s * paw.nkpts + k 

169 for n in range(paw.nbands): 

170 # avery: Added dummy loop body to make compiling work. 

171 1 

172'''