Coverage for gpaw/lcaotddft/densitymatrix.py: 90%

1import numpy as np 

2 

3from gpaw.utilities import pack_density 

4 

5 

6def get_density(rho_MM, wfs, density, density_type='comp', u=0): 

7 if wfs.ksl.using_blacs: 

8 raise NotImplementedError('Scalapack is not supported') 

9 

10 rho_G = density.gd.zeros() 

11 kpt = wfs.kpt_u[u] 

12 assert kpt.q == 0 

13 wfs.basis_functions.construct_density(rho_MM.astype(wfs.dtype), 

14 rho_G, kpt.q) 

15 

16 # Uncomment this if you want to add the static part 

17 # rho_G += density.nct_G 

18 

19 if density_type == 'pseudocoarse': 

20 return rho_G 

21 

22 rho_g = density.finegd.zeros() 

23 density.distribute_and_interpolate(rho_G, rho_g) 

24 rho_G = None 

25 

26 if density_type == 'pseudo': 

27 return rho_g 

28 

29 if density_type == 'comp': 

30 D_asp = density.atom_partition.arraydict(density.D_asp.shapes_a) 

31 Q_aL = {} 

32 for a, D_sp in D_asp.items(): 

33 P_Mi = wfs.P_aqMi[a][kpt.q] 

34 assert np.max(np.absolute(P_Mi.imag)) == 0 

35 P_Mi = P_Mi.real 

36 assert P_Mi.dtype == float 

37 D_ii = np.dot(np.dot(P_Mi.T.conj(), rho_MM), P_Mi) 

38 D_sp[:] = pack_density(D_ii)[np.newaxis, :] 

39 Q_aL[a] = np.dot(D_sp.sum(axis=0), wfs.setups[a].Delta_pL) 

40 density.ghat.add(rho_g, Q_aL) 

41 return rho_g 

42 

43 raise RuntimeError('Unknown density type: %s' % density_type) 

44 

45 

46class DensityMatrix: 

47 

48 def __init__(self, paw): 

49 self.wfs = paw.wfs 

50 self.density = paw.density 

51 self.using_blacs = self.wfs.ksl.using_blacs 

52 self.tag = None 

53 

54 def zeros(self, dtype): 

55 ksl = self.wfs.ksl 

56 if self.using_blacs: 

57 return ksl.mmdescriptor.zeros(dtype=dtype) 

58 else: 

59 return np.zeros((ksl.mynao, ksl.nao), dtype=dtype) 

60 

61 def _calculate_density_matrix(self, wfs, kpt): 

62 if self.using_blacs: 

63 ksl = wfs.ksl 

64 rho_MM = ksl.calculate_blocked_density_matrix(kpt.f_n, kpt.C_nM) 

65 else: 

66 rho_MM = wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM) 

67 wfs.bd.comm.sum(rho_MM, root=0) 

68 # TODO: should the sum over bands be moved to 

69 # OrbitalLayouts.calculate_density_matrix() 

70 return rho_MM 

71 

72 def get_density_matrix(self, tag=None): 

73 if tag is None or self.tag != tag: 

74 self.rho_uMM = [] 

75 for kpt in self.wfs.kpt_u: 

76 rho_MM = self._calculate_density_matrix(self.wfs, kpt) 

77 self.rho_uMM.append(rho_MM) 

78 self.tag = tag 

79 return self.rho_uMM 

80 

81 def get_density(self, rho_uMM=None, density_type='comp'): 

82 assert len(self.wfs.kpt_u) == 1, 'K-points not implemented' 

83 u = 0 

84 if rho_uMM is None: 

85 rho_uMM = self.rho_uMM 

86 return get_density(rho_uMM[u], self.wfs, self.density, density_type, u)