Coverage for gpaw/new/spinspiral.py: 33%

1from math import pi 

2from gpaw.core.plane_waves import PWDesc 

3from gpaw.core.pwacf import PWAtomCenteredFunctions 

4 

5 

6class SpiralPWDesc: 

7 def __init__(self, 

8 pw: PWDesc, 

9 qspiral_v): 

10 self.pw = pw 

11 self.qspiral_v = qspiral_v 

12 self.dtype = complex 

13 self.maxmysize = pw.maxmysize 

14 self.comm = pw.comm 

15 self.myshape = pw.myshape 

16 self.G_plus_k_Gv = pw.G_plus_k_Gv + qspiral_v 

17 self.ekin_G = 0.5 * (self.G_plus_k_Gv**2).sum(1) 

18 self.kpt = pw.kpt_c + pw.cell_cv @ qspiral_v / (2 * pi) 

19 self.kpt_c = self.kpt 

20 self.cell = pw.cell 

21 self.dv = pw.dv 

22 

23 

24class SpiralPWACF: 

25 def __init__(self, functions, positions, pw, 

26 atomdist, 

27 qspiral_v): 

28 self.pt_saiG = [ 

29 PWAtomCenteredFunctions( 

30 functions, 

31 positions, 

32 SpiralPWDesc(pw, 0.5 * sign * qspiral_v), 

33 atomdist=atomdist) 

34 for sign in [1, -1]] 

35 

36 def empty(self, dims, comm): 

37 return self.pt_saiG[0].empty(dims, comm) 

38 

39 def integrate(self, psit_nsG, out): 

40 P_ansi = out 

41 for s, pt_aiG in enumerate(self.pt_saiG): 

42 pt_aiG._lazy_init() 

43 pt_aiG._lfc.integrate(psit_nsG.data[:, s], 

44 {a: P_nsi[:, s] 

45 for a, P_nsi in P_ansi.items()}) 

46 

47 def add_to(self, r_nsG, P_ansi): 

48 for s, pt_aiG in enumerate(self.pt_saiG): 

49 # pt_aiG._lazy_init() 

50 pt_aiG._lfc.add(r_nsG.data[:, s], 

51 {a: P_nsi[:, s] 

52 for a, P_nsi in P_ansi.items()})