Coverage for gpaw/hybrids/wstc.py: 72%

1"""Wigner-Seitz truncated coulomb interaction. 

2 

3See: 

4 

5 Ravishankar Sundararaman and T. A. Arias: 

6 Phys. Rev. B 87, 165122 (2013) 

7 

8 Regularization of the Coulomb singularity in exact exchange by 

9 Wigner-Seitz truncated interactions: Towards chemical accuracy 

10 in nontrivial systems 

11""" 

12from math import pi 

13 

14import numpy as np 

15from scipy.special import erf 

16from ase.units import Bohr 

17from ase.utils import seterr 

18 

19import gpaw.mpi as mpi 

20from gpaw.fftw import get_efficient_fft_size 

21from gpaw.grid_descriptor import GridDescriptor 

22from gpaw.core import PWDesc, UGDesc, PWArray 

23from gpaw.typing import Array1D 

24 

25 

26class WignerSeitzTruncatedCoulomb: 

27 def __init__(self, cell_cv, nk_c): 

28 self.nk_c = nk_c 

29 bigcell_cv = cell_cv * nk_c[:, np.newaxis] 

30 L_c = (np.linalg.inv(bigcell_cv)**2).sum(0)**-0.5 

31 

32 self.rc = 0.5 * L_c.min() 

33 self.a = 5 / self.rc 

34 

35 nr_c = [get_efficient_fft_size(2 * int(L * self.a * 3.0)) 

36 for L in L_c] 

37 

38 self.gd = GridDescriptor(nr_c, bigcell_cv, comm=mpi.serial_comm) 

39 v_ijk = self.gd.empty() 

40 

41 pos_ijkv = self.gd.get_grid_point_coordinates().transpose((1, 2, 3, 0)) 

42 corner_xv = np.dot(np.indices((2, 2, 2)).reshape((3, 8)).T, bigcell_cv) 

43 

44 # Ignore division by zero (in 0,0,0 corner): 

45 with seterr(invalid='ignore'): 

46 # Loop over first dimension to avoid too large ndarrays. 

47 for pos_jkv, v_jk in zip(pos_ijkv, v_ijk): 

48 # Distances to the 8 corners: 

49 d_jkxv = pos_jkv[:, :, np.newaxis] - corner_xv 

50 r_jk = (d_jkxv**2).sum(axis=3).min(2)**0.5 

51 v_jk[:] = erf(self.a * r_jk) / r_jk 

52 

53 # Fix 0/0 corner value: 

54 v_ijk[0, 0, 0] = 2 * self.a / pi**0.5 

55 

56 self.K_Q = np.fft.fftn(v_ijk) * self.gd.dv 

57 

58 def get_description(self): 

59 descriptors = [] 

60 descriptors.append('Inner radius for %dx%dx%d Wigner-Seitz cell: ' 

61 '%.3f Ang' % (tuple(self.nk_c) + (self.rc * Bohr,))) 

62 descriptors.append('Range-separation parameter: %.3f Ang^-1' % ( 

63 self.a / Bohr)) 

64 descriptors.append('FFT size for calculating truncated Coulomb: ' 

65 '%dx%dx%d' % tuple(self.gd.N_c)) 

66 

67 return '\n'.join(descriptors) 

68 

69 def get_potential_new(self, pw: PWDesc, grid: UGDesc) -> PWArray: 

70 K_G = pw.zeros() 

71 assert isinstance(K_G, PWArray) # !!! 

72 size = tuple(grid.size) 

73 if pw.dtype == float: 

74 size = (size[0], size[1], size[2] // 2 + 1) 

75 self._get_potential(pw.indices(size)[pw.ng1:pw.ng2], 

76 pw.kpt_c, 

77 grid.size_c, 

78 pw.dtype, 

79 pw.reciprocal_vectors(), 

80 K_G.data) 

81 return K_G 

82 

83 def get_potential(self, pd) -> Array1D: 

84 K_G = pd.zeros() 

85 self._get_potential(pd.Q_qG[0], 

86 pd.kd.bzk_kc[0], 

87 pd.gd.N_c, 

88 pd.dtype, 

89 pd.get_reciprocal_vectors(add_q=True), 

90 K_G) 

91 return K_G 

92 

93 def _get_potential(self, Q_G, q_c, N_c, dtype, qG_Gv, K_G) -> None: 

94 shift_c = (q_c * self.nk_c).round().astype(int) 

95 max_c = self.gd.N_c // 2 

96 if dtype == complex: 

97 for G, Q in enumerate(Q_G): 

98 Q_c = (np.unravel_index(Q, N_c) + N_c // 2) % N_c - N_c // 2 

99 Q_c = Q_c * self.nk_c + shift_c 

100 if (abs(Q_c) < max_c).all(): 

101 K_G[G] = self.K_Q[tuple(Q_c)] 

102 else: 

103 Nr_c = N_c.copy() 

104 Nr_c[2] = N_c[2] // 2 + 1 

105 for G, Q in enumerate(Q_G): 

106 Q_c = np.array(np.unravel_index(Q, Nr_c)) 

107 Q_c[:2] += N_c[:2] // 2 

108 Q_c[:2] %= N_c[:2] 

109 Q_c[:2] -= N_c[:2] // 2 

110 if (abs(Q_c) < max_c).all(): 

111 K_G[G] = self.K_Q[tuple(Q_c)] 

112 

113 G2_G = np.sum(qG_Gv**2, axis=1) 

114 # G2_G = pd.G2_qG[0] 

115 a = self.a 

116 G0 = G2_G.argmin() 

117 if G2_G[G0] < 1e-11: 

118 K0 = K_G[G0] + pi / a**2 

119 with np.errstate(invalid='ignore'): 

120 K_G += 4 * pi * (1 - np.exp(-G2_G / (4 * a**2))) / G2_G 

121 if G2_G[G0] < 1e-11: 

122 K_G[G0] = K0