1import re
2import numpy as np
3
4from gpaw import __version__ as version
5from gpaw.mpi import world
6from gpaw.tddft.units import au_to_as, au_to_fs, au_to_eV, rot_au_to_cgs
7from gpaw.tddft.folding import FoldedFrequencies
8from gpaw.tddft.folding import Folding
9
10
11def calculate_fourier_transform(x_t, y_ti, foldedfrequencies):
12    ff = foldedfrequencies
13    X_w = ff.frequencies
14    envelope = ff.folding.envelope
15
16    # Construct integration weights:
17    # We use trapezoidal rule except the end point is accounted with
18    # full weight. This ensures better numerical compatibility
19    # with the on-the-fly Fourier integrators.
20    # This shouldn't affect in usual scenarios as the envelope
21    # should damp the data to zero at the end point in any case.
22    dx_t1 = x_t[1:] - x_t[:-1]
23    dx_t = 0.5 * (np.insert(dx_t1, 0, 0.0) + np.append(dx_t1, dx_t1[-1]))
24
25    # Integrate
26    f_wt = np.exp(1.0j * np.outer(X_w, x_t))
27    y_it = np.swapaxes(y_ti, 0, 1)
28    Y_wi = np.tensordot(f_wt, dx_t * envelope(x_t) * y_it, axes=(1, 1))
29    return Y_wi
30
31
32def read_td_file_data(fname, remove_duplicates=True):
33    """Read data from time-dependent data file.
34
35    Parameters
36    ----------
37    fname
38        File path
39    remove_duplicates
40        If true, remove data from overlapping time values.
41        The first encountered values are kept.
42
43    Returns
44    -------
45    time_t
46        Array of time values
47    data_ti
48        Array of data values
49    """
50    # Read data
51    data_tj = np.loadtxt(fname)
52    time_t = data_tj[:, 0]
53    data_ti = data_tj[:, 1:]
54
55    # Remove duplicates due to abruptly stopped and restarted calculation
56    if remove_duplicates:
57        flt_t = np.ones_like(time_t, dtype=bool)
58        maxtime = time_t[0]
59        for t in range(1, len(time_t)):
60            # Note about ">=" here:
61            # The equality is included here in order to
62            # retain step-like data (for example, the data both just before
63            # and just after the kick is kept).
64            if time_t[t] >= maxtime:
65                maxtime = time_t[t]
66            else:
67                flt_t[t] = False
68        time_t = time_t[flt_t]
69        data_ti = data_ti[flt_t]
70        ndup = len(flt_t) - flt_t.sum()
71        if ndup > 0:
72            print('Removed %d duplicates' % ndup)
73    return time_t, data_ti
74
75
76def read_td_file_kicks(fname):
77    """Read kicks from time-dependent data file.
78
79    Parameters
80    ----------
81    fname
82        File path
83
84    Returns
85    -------
86    kick_i
87        List of kicks.
88        Each kick is a dictionary with keys
89        ``strength_v`` and ``time``.
90    """
91    def parse_kick_line(line):
92        # Kick
93        regexp = (r"Kick = \["
94                  r"(?P<k0>[-+0-9\.e\ ]+), "
95                  r"(?P<k1>[-+0-9\.e\ ]+), "
96                  r"(?P<k2>[-+0-9\.e\ ]+)\]")
97        m = re.search(regexp, line)
98        assert m is not None, 'Kick not found'
99        kick_v = np.array([float(m.group('k%d' % v)) for v in range(3)])
100        # Time
101        regexp = r"Time = (?P<time>[-+0-9\.e\ ]+)"
102        m = re.search(regexp, line)
103        if m is None:
104            print('time not found')
105            time = 0.0
106        else:
107            time = float(m.group('time'))
108        return kick_v, time
109
110    # Search kicks
111    kick_i = []
112    with open(fname, 'r') as f:
113        for line in f:
114            if line.startswith('# Kick'):
115                kick_v, time = parse_kick_line(line)
116                kick_i.append({'strength_v': kick_v, 'time': time})
117    return kick_i
118
119
120def clean_td_data(kick_i, time_t, data_ti):
121    """Prune time-dependent data.
122
123    This function checks that there is only one kick
124    in the kick list and moves time zero to the kick
125    time (discarding all preceding data).
126
127    Parameters
128    ----------
129    kick_i
130        List of kicks.
131    time_t
132        Array of time values
133    data_ti
134        Array of data values
135
136    Returns
137    -------
138    kick_i
139        List of kicks.
140        Each kick is a dictionary with keys
141        ``strength_v`` and ``time``.
142
143    Raises
144    ------
145    RuntimeError
146        If kick list contains multiple kicks.
147    """
148    # Check kicks
149    if len(kick_i) > 1:
150        raise RuntimeError('Multiple kicks')
151    kick = kick_i[0]
152    kick_v = kick['strength_v']
153    kick_time = kick['time']
154
155    # Discard times before kick
156    flt_t = time_t >= kick_time
157    time_t = time_t[flt_t]
158    data_ti = data_ti[flt_t]
159
160    # Move time zero to kick time
161    time_t -= kick_time
162    assert time_t[0] == 0.0
163
164    return kick_v, time_t, data_ti
165
166
167def read_dipole_moment_file(fname, remove_duplicates=True):
168    """Read time-dependent dipole moment data file.
169
170    Parameters
171    ----------
172    fname
173        File path
174    remove_duplicates
175        If true, remove data from overlapping time values.
176        The first encountered values are kept.
177
178    Returns
179    -------
180    kick_i
181        List of kicks.
182        Each kick is a dictionary with keys
183        ``strength_v`` and ``time``.
184    time_t
185        Array of time values
186    norm_t
187        Array of norm values
188    dm_tv
189        Array of dipole moment values
190    """
191    time_t, data_ti = read_td_file_data(fname, remove_duplicates)
192    kick_i = read_td_file_kicks(fname)
193    norm_t = data_ti[:, 0]
194    dm_tv = data_ti[:, 1:]
195    return kick_i, time_t, norm_t, dm_tv
196
197
198def calculate_polarizability(kick_v, time_t, dm_tv, foldedfrequencies):
199    dm_tv = dm_tv - dm_tv[0]
200    alpha_wv = calculate_fourier_transform(time_t, dm_tv, foldedfrequencies)
201    kick_magnitude = np.sqrt(np.sum(kick_v**2))
202    alpha_wv /= kick_magnitude
203    return alpha_wv
204
205
206def calculate_photoabsorption(kick_v, time_t, dm_tv, foldedfrequencies):
207    omega_w = foldedfrequencies.frequencies
208    alpha_wv = calculate_polarizability(kick_v, time_t, dm_tv,
209                                        foldedfrequencies)
210    abs_wv = 2 / np.pi * omega_w[:, np.newaxis] * alpha_wv.imag
211
212    kick_magnitude = np.sqrt(np.sum(kick_v**2))
213    abs_wv *= kick_v / kick_magnitude
214    return abs_wv
215
216
217def read_magnetic_moment_file(fname, remove_duplicates=True):
218    """Read time-dependent magnetic moment data file.
219
220    Parameters
221    ----------
222    fname
223        File path
224    remove_duplicates
225        If true, remove data from overlapping time values.
226        The first encountered values are kept.
227
228    Returns
229    -------
230    kick_i
231        List of kicks.
232        Each kick is a dictionary with keys
233        ``strength_v`` and ``time``.
234    time_t
235        Array of time values
236    mm_tv
237        Array of magnetic moment values
238    """
239    time_t, mm_tv = read_td_file_data(fname, remove_duplicates)
240    kick_i = read_td_file_kicks(fname)
241    return kick_i, time_t, mm_tv
242
243
244def calculate_rotatory_strength_components(kick_v, time_t, mm_tv,
245                                           foldedfrequencies):
246    assert np.all(mm_tv[0] == 0.0)
247    mm_wv = calculate_fourier_transform(time_t, mm_tv, foldedfrequencies)
248    kick_magnitude = np.sqrt(np.sum(kick_v**2))
249    rot_wv = mm_wv.real / (np.pi * kick_magnitude)
250    return rot_wv
251
252
253def write_spectrum(dipole_moment_file, spectrum_file,
254                   folding, width, e_min, e_max, delta_e,
255                   title, symbol, calculate):
256    def str_list(v_i, fmt='%g'):
257        return '[%s]' % ', '.join(map(lambda v: fmt % v, v_i))
258
259    kick_i, time_t, _, dm_tv = read_dipole_moment_file(dipole_moment_file)
260    kick_v, time_t, dm_tv = clean_td_data(kick_i, time_t, dm_tv)
261    dt_t = time_t[1:] - time_t[:-1]
262
263    freqs = np.arange(e_min, e_max + 0.5 * delta_e, delta_e)
264    folding = Folding(folding, width)
265    ff = FoldedFrequencies(freqs, folding)
266    omega_w = ff.frequencies
267    spec_wv = calculate(kick_v, time_t, dm_tv, ff)
268
269    # Write spectrum file header
270    with open(spectrum_file, 'w') as f:
271        def w(s):
272            f.write('%s\n' % s)
273
274        w('# %s spectrum from real-time propagation' % title)
275        w('# GPAW version: %s' % version)
276        w('# Total time = %.4f fs, Time steps = %s as' %
277          (dt_t.sum() * au_to_fs,
278           str_list(np.unique(np.around(dt_t, 6)) * au_to_as, '%.4f')))
279        w('# Kick = %s' % str_list(kick_v))
280        w('# %sian folding, Width = %.4f eV = %lf Hartree'
281          ' <=> FWHM = %lf eV' %
282          (folding.folding, folding.width * au_to_eV, folding.width,
283           folding.fwhm * au_to_eV))
284
285        col_i = []
286        data_iw = [omega_w * au_to_eV]
287        for v in range(len(kick_v)):
288            h = '%s_%s' % (symbol, 'xyz'[v])
289            if spec_wv.dtype == complex:
290                col_i.append('Re[%s]' % h)
291                data_iw.append(spec_wv[:, v].real)
292                col_i.append('Im[%s]' % h)
293                data_iw.append(spec_wv[:, v].imag)
294            else:
295                col_i.append(h)
296                data_iw.append(spec_wv[:, v])
297
298        w('# %10s %s' % ('om (eV)', ' '.join(['%20s' % s for s in col_i])))
299
300    # Write spectrum file data
301    with open(spectrum_file, 'ab') as f:
302        np.savetxt(f, np.array(data_iw).T,
303                   fmt='%12.6lf' + (' %20.10le' * len(col_i)))
304
305    return folding.envelope(time_t[-1])
306
307
308def photoabsorption_spectrum(dipole_moment_file: str,
309                             spectrum_file: str,
310                             folding: str = 'Gauss',
311                             width: float = 0.2123,
312                             e_min: float = 0.0,
313                             e_max: float = 30.0,
314                             delta_e: float = 0.05):
315    """Calculates photoabsorption spectrum from the time-dependent
316    dipole moment.
317
318    Parameters
319    ----------
320    dipole_moment_file
321        Name of the time-dependent dipole moment file from which
322        the spectrum is calculated
323    spectrum_file
324        Name of the spectrum file
325    folding
326        Gaussian (``'Gauss'``) or Lorentzian (``'Lorentz'``) folding
327    width
328        Width of the Gaussian (sigma) or Lorentzian (Gamma)
329        Gaussian =     1/(sigma sqrt(2pi)) exp(-(1/2)(omega/sigma)^2)
330        Lorentzian =  (1/pi) (1/2) Gamma / [omega^2 + ((1/2) Gamma)^2]
331    e_min
332        Minimum energy shown in the spectrum (eV)
333    e_max
334        Maximum energy shown in the spectrum (eV)
335    delta_e
336        Energy resolution (eV)
337    """
338    if world.rank == 0:
339        print('Calculating photoabsorption spectrum from file "%s"'
340              % dipole_moment_file)
341
342        def calculate(*args):
343            return calculate_photoabsorption(*args) / au_to_eV
344        sinc = write_spectrum(dipole_moment_file, spectrum_file,
345                              folding, width, e_min, e_max, delta_e,
346                              'Photoabsorption', 'S', calculate)
347        print('Sinc contamination %.8f' % sinc)
348        print('Calculated photoabsorption spectrum saved to file "%s"'
349              % spectrum_file)
350
351
352def polarizability_spectrum(dipole_moment_file, spectrum_file,
353                            folding='Gauss', width=0.2123,
354                            e_min=0.0, e_max=30.0, delta_e=0.05):
355    """Calculates polarizability spectrum from the time-dependent
356    dipole moment.
357
358    Parameters:
359
360    dipole_moment_file: string
361        Name of the time-dependent dipole moment file from which
362        the spectrum is calculated
363    spectrum_file: string
364        Name of the spectrum file
365    folding: 'Gauss' or 'Lorentz'
366        Whether to use Gaussian or Lorentzian folding
367    width: float
368        Width of the Gaussian (sigma) or Lorentzian (Gamma)
369        Gaussian =     1/(sigma sqrt(2pi)) exp(-(1/2)(omega/sigma)^2)
370        Lorentzian =  (1/pi) (1/2) Gamma / [omega^2 + ((1/2) Gamma)^2]
371    e_min: float
372        Minimum energy shown in the spectrum (eV)
373    e_max: float
374        Maximum energy shown in the spectrum (eV)
375    delta_e: float
376        Energy resolution (eV)
377    """
378    if world.rank == 0:
379        print('Calculating polarizability spectrum from file "%s"'
380              % dipole_moment_file)
381
382        def calculate(*args):
383            return calculate_polarizability(*args) / au_to_eV**2
384        sinc = write_spectrum(dipole_moment_file, spectrum_file,
385                              folding, width, e_min, e_max, delta_e,
386                              'Polarizability', 'alpha', calculate)
387        print('Sinc contamination %.8f' % sinc)
388        print('Calculated polarizability spectrum saved to file "%s"'
389              % spectrum_file)
390
391
392def rotatory_strength_spectrum(magnetic_moment_files, spectrum_file,
393                               folding='Gauss', width=0.2123,
394                               e_min=0.0, e_max=30.0, delta_e=0.05):
395    """Calculates rotatory strength spectrum from the time-dependent
396    magnetic moment.
397
398    Parameters
399    ----------
400    magnetic_moment_files: list of string
401        Time-dependent magnetic moment files for x, y, and z kicks
402    spectrum_file: string
403        Name of the spectrum file
404    folding: 'Gauss' or 'Lorentz'
405        Whether to use Gaussian or Lorentzian folding
406    width: float
407        Width of the Gaussian (sigma) or Lorentzian (Gamma)
408        Gaussian =     1/(sigma sqrt(2pi)) exp(-(1/2)(omega/sigma)^2)
409        Lorentzian =  (1/pi) (1/2) Gamma / [omega^2 + ((1/2) Gamma)^2]
410    e_min: float
411        Minimum energy shown in the spectrum (eV)
412    e_max: float
413        Maximum energy shown in the spectrum (eV)
414    delta_e: float
415        Energy resolution (eV)
416    """
417    if world.rank != 0:
418        return
419
420    freqs = np.arange(e_min, e_max + 0.5 * delta_e, delta_e)
421    folding = Folding(folding, width)
422    ff = FoldedFrequencies(freqs, folding)
423    omega_w = ff.frequencies * au_to_eV
424    rot_w = np.zeros_like(omega_w)
425
426    tot_time = np.inf
427    time_steps = []
428    kick_strength = None
429    for v, fpath in enumerate(magnetic_moment_files):
430        kick_i, time_t, mm_tv = read_magnetic_moment_file(fpath)
431        kick_v, time_t, mm_tv = clean_td_data(kick_i, time_t, mm_tv)
432
433        tot_time = min(tot_time, time_t[-1])
434        time_steps.append(np.around(time_t[1:] - time_t[:-1], 6))
435
436        # Check kicks
437        for v0 in range(3):
438            if v0 == v:
439                continue
440            if kick_v[v0] != 0.0:
441                raise RuntimeError('The magnetic moment files must be '
442                                   'for kicks in x, y, and z directions.')
443        if kick_strength is None:
444            kick_strength = np.sqrt(np.sum(kick_v**2))
445        if np.sqrt(np.sum(kick_v**2)) != kick_strength:
446            raise RuntimeError('The magnetic moment files must have '
447                               'been calculated with the same kick strength.')
448
449        rot_wv = calculate_rotatory_strength_components(kick_v, time_t,
450                                                        mm_tv, ff)
451        rot_w += rot_wv[:, v]
452
453    rot_w *= rot_au_to_cgs * 1e40 / au_to_eV
454
455    # Unique non-zero time steps
456    time_steps = np.unique(time_steps)
457    time_steps = time_steps[time_steps != 0]
458
459    with open(spectrum_file, 'w') as fd:
460        steps_str = ', '.join(f'{val:.4f}' for val in time_steps * au_to_as)
461        lines = ['Rotatory strength spectrum from real-time propagations',
462                 f'Total time = {tot_time * au_to_fs:.4f} fs, '
463                 f'Time steps = [{steps_str}] as',
464                 f'Kick strength = {kick_strength}',
465                 f'{folding.folding}ian folding, '
466                 f'width = {folding.width * au_to_eV:.4f} eV '
467                 f'<=> FWHM = {folding.fwhm * au_to_eV:.4f} eV']
468        fd.write('# ' + '\n# '.join(lines) + '\n')
469        fd.write(f'# {"Energy (eV)":>12} {"R (1e-40 cgs / eV)":>20}\n')
470
471        data_wi = np.vstack((omega_w, rot_w)).T
472        np.savetxt(fd, data_wi,
473                   fmt='%14.6lf' + (' %20.10le' * (data_wi.shape[1] - 1)))
474