Coverage for gpaw/lcaotddft/dipolemomentwriter.py: 83%
162 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1import re
3import numpy as np
5from ase.utils import IOContext
7from gpaw.lcaotddft.magneticmomentwriter import parse_header
9from gpaw.lcaotddft.observer import TDDFTObserver
12def convert_repr(r):
13 # Integer
14 try:
15 return int(r)
16 except ValueError:
17 pass
18 # Boolean
19 b = {repr(False): False, repr(True): True}.get(r, None)
20 if b is not None:
21 return b
22 # String
23 s = r[1:-1]
24 if repr(s) == r:
25 return s
26 raise RuntimeError('Unknown value: %s' % r)
29class DipoleMomentWriter(TDDFTObserver):
30 """Observer for writing time-dependent dipole moment data.
32 The data is written in atomic units.
34 The observer attaches to the TDDFT calculator during creation.
36 Parameters
37 ----------
38 paw
39 TDDFT calculator
40 filename
41 File for writing dipole moment data
42 center
43 If true, dipole moment is evaluated with the center of cell
44 as the origin
45 density
46 Density type used for evaluating dipole moment.
47 Use the default value for production calculations;
48 others are for testing purposes.
49 Possible values:
50 ``'comp'``: ``rhot_g``,
51 ``'pseudo'``: ``nt_sg``,
52 ``'pseudocoarse'``: ``nt_sG``.
53 force_new_file
54 If true, new dipole moment file is created (erasing any existing one)
55 even when restarting time propagation.
56 interval
57 Update interval. Value of 1 corresponds to evaluating and
58 writing data after every propagation step.
59 """
60 version = 1
62 def __init__(self, paw, filename: str, *,
63 center: bool = False,
64 density: str = 'comp',
65 force_new_file: bool = False,
66 interval: int = 1):
67 TDDFTObserver.__init__(self, paw, interval)
68 self.ioctx = IOContext()
69 if paw.niter == 0 or force_new_file:
70 # Initialize
71 self.do_center = center
72 self.density_type = density
73 self.fd = self.ioctx.openfile(filename, comm=paw.world, mode='w')
74 self._write_header(paw)
75 else:
76 # Read and continue
77 self.read_header(filename)
78 self.fd = self.ioctx.openfile(filename, comm=paw.world, mode='a')
80 def _write(self, line):
81 self.fd.write(line)
82 self.fd.flush()
84 def _write_header(self, paw):
85 line = f'# {self.__class__.__name__}[version={self.version}]'
86 line += ('(center=%s, density=%s)\n' %
87 (repr(self.do_center), repr(self.density_type)))
88 line += ('# %15s %15s %22s %22s %22s\n' %
89 ('time', 'norm', 'dmx', 'dmy', 'dmz'))
90 self._write(line)
92 def read_header(self, filename):
93 with open(filename) as f:
94 line = f.readline()
95 m_i = re.split("[^a-zA-Z0-9_=']+", line[2:])
96 assert m_i.pop(0) == self.__class__.__name__
97 for m in m_i:
98 if '=' not in m:
99 continue
100 k, v = m.split('=')
101 v = convert_repr(v)
102 if k == 'version':
103 assert v == self.version
104 continue
105 # Translate key
106 k = {'center': 'do_center', 'density': 'density_type'}[k]
107 setattr(self, k, v)
109 def _write_init(self, paw):
110 time = paw.time
111 line = '# Start; Time = %.8lf\n' % time
112 self._write(line)
114 def _write_kick(self, paw):
115 time = paw.time
116 kick = paw.kick_strength
117 gauge = paw.kick_gauge
118 line = '# Kick = [%22.12le, %22.12le, %22.12le]; ' % tuple(kick)
119 line += 'Gauge = %s; ' % gauge
120 line += 'Time = %.8lf\n' % time
121 self._write(line)
123 def calculate_dipole_moment(self, gd, rho_g, center=True):
124 if center:
125 center_v = 0.5 * gd.cell_cv.sum(0)
126 else:
127 center_v = np.zeros(3, dtype=float)
128 r_vg = gd.get_grid_point_coordinates()
129 dm_v = np.zeros(3, dtype=float)
130 for v in range(3):
131 dm_v[v] = - gd.integrate((r_vg[v] - center_v[v]) * rho_g)
132 return dm_v
134 def _write_dm(self, paw):
135 time = paw.time
136 density = paw.density
137 if self.density_type == 'comp':
138 rho_g = density.rhot_g
139 gd = density.finegd
140 elif self.density_type == 'pseudo':
141 rho_g = density.nt_sg.sum(axis=0)
142 gd = density.finegd
143 elif self.density_type == 'pseudocoarse':
144 rho_g = density.nt_sG.sum(axis=0)
145 gd = density.gd
146 else:
147 raise RuntimeError('Unknown density type: %s' % self.density_type)
149 norm = gd.integrate(rho_g)
150 # dm = self.calculate_dipole_moment(gd, rho_g, center=self.do_center)
151 dm = gd.calculate_dipole_moment(rho_g, center=self.do_center)
152 if paw.hamiltonian.poisson.get_description() == 'FDTD+TDDFT': # XXX
153 dm += paw.hamiltonian.poisson.get_classical_dipole_moment()
154 line = ('%20.8lf %20.8le %22.12le %22.12le %22.12le\n' %
155 (time, norm, dm[0], dm[1], dm[2]))
156 self._write(line)
158 def _update(self, paw):
159 if paw.action == 'init':
160 self._write_init(paw)
161 elif paw.action == 'kick':
162 self._write_kick(paw)
163 self._write_dm(paw)
165 def __del__(self):
166 self.ioctx.close()
169# TODO: Create a common base class for VelocityGaugeWriter
170# and DipoleMomentWriter.
171class VelocityGaugeWriter(TDDFTObserver):
172 """Observer for writing time-dependent velocity-kick density
173 for test purposes.
175 The data is written in atomic units.
177 The observer attaches to the TDDFT calculator during creation.
179 Parameters
180 ----------
181 paw
182 TDDFT calculator
183 filename
184 File for writing density data
185 """
186 version = 5
188 def __init__(self, paw, filename: str, interval: int = 1):
189 super().__init__(paw, interval)
190 self.ioctx = IOContext()
191 mode = paw.wfs.mode
192 assert mode in ['lcao']
194 if paw.niter == 0:
195 self.fd = self.ioctx.openfile(filename, comm=paw.world, mode='w')
196 else:
197 self.fd = self.ioctx.openfile(filename, comm=paw.world, mode='a')
199 def _write(self, line):
200 self.fd.write(line)
201 self.fd.flush()
203 def _write_header(self, paw, kwargs):
204 line = f'# {self.__class__.__name__}[version={self.version}]\n'
205 line += ('# %15s %15s %22s %22s %22s\n' %
206 ('time', 'norm', 'rhoVMM_x', 'rhoVMM_y', 'rhoVMM_z'))
207 self._write(line)
209 def _read_header(self, filename):
210 with open(filename, encoding='utf-8') as fd:
211 line = fd.readline()
212 try:
213 name, version, kwargs = parse_header(line[2:])
214 except ValueError as e:
215 raise ValueError(f'File {filename} cannot be parsed: {e}')
216 if name != self.__class__.__name__:
217 raise ValueError(f'File {filename} is not '
218 f'for {self.__class__.__name__}')
219 if version != self.version:
220 raise ValueError(f'File {filename} is not '
221 f'of version {self.version}')
222 return kwargs
224 def _write_init(self, paw):
225 time = paw.time
226 line = '# Start; Time = %.8lf\n' % time
227 self._write(line)
229 def _write_kick(self, paw):
230 time = paw.time
231 kick = paw.kick_strength
232 gauge = paw.kick_gauge
233 line = '# Kick = [%22.12le, %22.12le, %22.12le]; ' % tuple(kick)
234 line += 'Gauge = %s; ' % gauge
235 line += 'Time = %.8lf\n' % time
236 self._write(line)
238 def _calculate_v(self, paw):
239 C_nM = paw.wfs.kpt_u[0].C_nM
240 f_n = paw.wfs.kpt_u[0].f_n
241 rho_MM = C_nM.T.conj() @ (f_n[:, None] * C_nM)
242 Vkick_vmm = paw.wfs.kpt_u[0].Vkick_vmm
243 return [np.sum((rho_MM * Vkick_mm.conj()).ravel())
244 for Vkick_mm in Vkick_vmm]
246 def _write_v(self, paw):
247 time = paw.time
248 v_v = [v.real for v in self._calculate_v(paw)]
249 line = ('%20.8lf %22.12le %22.12le %22.12le %22.12le\n'
250 % (time, 0.0, v_v[0], v_v[1], v_v[2]))
251 self._write(line)
253 def _update(self, paw): # This does not seem to work in here
254 if paw.action == 'init':
255 paw.propagator.calculate_velocity_operator_matrix()
256 self._write_header(paw, {})
257 self._write_init(paw)
258 elif paw.action == 'kick':
259 self._write_kick(paw)
260 self._write_v(paw)
262 def __del__(self):
263 self.ioctx.close()