Coverage for gpaw/lcaotddft/quadrupolemomentwriter.py: 19%

1import re 

2import numpy as np 

3 

4from ase.utils import IOContext 

5 

6from gpaw.lcaotddft.observer import TDDFTObserver 

7 

8 

9def calculate_quadrupole_moments(gd, rho_g, center_v): 

10 # Center relative to the cell center 

11 # center_v += 0.5 * gd.cell_cv.sum(0) 

12 r_vg = gd.get_grid_point_coordinates() 

13 qm_i = np.zeros(6, dtype=float) 

14 i = 0 

15 for v1 in range(3): 

16 x1_g = r_vg[v1] - center_v[v1] 

17 for v2 in range(v1, 3): 

18 x2_g = r_vg[v2] - center_v[v2] 

19 qm_i[i] = -gd.integrate(x1_g * x2_g * rho_g) 

20 i += 1 

21 return qm_i 

22 

23 

24class QuadrupoleMomentWriter(TDDFTObserver): 

25 version = 1 

26 

27 def __init__(self, paw, filename, center=[0, 0, 0], density='comp', 

28 interval=1): 

29 TDDFTObserver.__init__(self, paw, interval) 

30 self.ioctx = IOContext() 

31 if paw.niter == 0: 

32 # Initialize 

33 if isinstance(center, str) and center == 'center': 

34 self.center_v = 0.5 * paw.density.finegd.cell_cv.sum(0) 

35 else: 

36 assert len(center) == 3 

37 self.center_v = center 

38 self.density_type = density 

39 self.fd = self.ioctx.openfile(filename, comm=paw.world, mode='w') 

40 else: 

41 # Read and continue 

42 self.read_header(filename) 

43 self.fd = self.ioctx.openfile(filename, comm=paw.world, mode='a') 

44 

45 def _write(self, line): 

46 self.fd.write(line) 

47 self.fd.flush() 

48 

49 def _write_header(self, paw): 

50 if paw.niter != 0: 

51 return 

52 line = f'# {self.__class__.__name__}[version={self.version}]' 

53 line += '(center=[%.6f, %.6f, %.6f]' % tuple(self.center_v) 

54 line += ', density=%s)\n' % repr(self.density_type) 

55 line += ('# %15s %15s %22s %22s %22s %22s %22s %22s\n' % 

56 ('time', 'norm', 'x^2', 'xy', 'xz', 'y^2', 'yz', 'z^2')) 

57 self._write(line) 

58 

59 def read_header(self, filename): 

60 with open(filename) as f: 

61 line = f.readline() 

62 regexp = (r"^(?P<name>\w+)\[version=(?P<version>\d+)\]\(center=\[" 

63 r"(?P<center0>[-+0-9\.]+), " 

64 r"(?P<center1>[-+0-9\.]+), " 

65 r"(?P<center2>[-+0-9\.]+)\], " 

66 r"density='(?P<density>\w+)'\)$") 

67 m = re.match(regexp, line[2:]) 

68 assert m is not None, 'Unknown fileformat' 

69 assert m.group('name') == self.__class__.__name__ 

70 assert int(m.group('version')) == self.version 

71 self.center_v = [float(m.group('center%d' % v)) for v in range(3)] 

72 self.density_type = m.group('density') 

73 

74 def _write_kick(self, paw): 

75 line = '# Kick: %s' % paw.kick_ext 

76 line += '; Time = %.8lf' % paw.time 

77 line += '\n' 

78 self._write(line) 

79 

80 def _write_dm(self, paw): 

81 time = paw.time 

82 density = paw.density 

83 if self.density_type == 'comp': 

84 rho_g = density.rhot_g 

85 gd = density.finegd 

86 elif self.density_type == 'pseudo': 

87 rho_g = density.nt_sg.sum(axis=0) 

88 gd = density.finegd 

89 elif self.density_type == 'pseudocoarse': 

90 rho_g = density.nt_sG.sum(axis=0) 

91 gd = density.gd 

92 else: 

93 raise RuntimeError('Unknown density type: %s' % self.density_type) 

94 

95 norm = gd.integrate(rho_g) 

96 qm = calculate_quadrupole_moments(gd, rho_g, center_v=self.center_v) 

97 line = (('%20.8lf %20.8le' + ' %22.12le' * len(qm) + '\n') % 

98 ((time, norm) + tuple(qm))) 

99 self._write(line) 

100 

101 def _update(self, paw): 

102 if paw.action == 'init': 

103 self._write_header(paw) 

104 elif paw.action == 'kick': 

105 self._write_kick(paw) 

106 self._write_dm(paw) 

107 

108 def __del__(self): 

109 self.ioctx.close()