Coverage for strongcoca/response/gpaw.py: 100%
98 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-26 18:44 +0000
1# This module contains components adapted from GPAW:
2# gpaw.tddft.spectrum
3# https://gitlab.com/gpaw/gpaw/-/blob/aca9ed6f520d6a855013247119b50c630f5eebc9/gpaw/tddft/spectrum.py
5import re
6from collections import namedtuple
7from functools import partial
8from pathlib import Path
9from typing import List, Optional, Tuple, Union
10import numpy as np
11from ase import Atoms
12import ase.io
13from . import TimeDomainResponse
14from .utilities import Broadening
15from ..types import Path_str
17Kick = namedtuple('Kick', ['time', 'strength_v'])
20def read_gpaw_tddft(dipole_files: Union[Path_str, List[Path_str]],
21 broadening: Broadening,
22 atoms_file: Optional[Path_str] = None,
23 name: str = 'GPAWTDDFT') -> TimeDomainResponse:
24 """Builds a :class:`~strongcoca.response.TimeDomainResponse` instance
25 using time-dependent dipole moment data from GPAW real-time
26 time-dependent density functional theory (RT-TDDFT) calculations.
28 Parameters
29 ----------
30 dipole_files
31 Name(s) of file(s) with time-dependent dipole moment data; each file
32 is assumed to provide a different Cartesian direction;
33 if there is only one file, the response is assumed to be diagonal.
34 broadening
35 Artificial broadening applied to the response.
36 atoms_file
37 Name of file from which to read atomic coordinates; needs to be readable
38 by `ASE <https://ase-lib.org/ase/io/io.html>`_.
39 name
40 Name of the response.
42 Returns
43 -------
44 response
45 Time domain response.
47 Raises
48 ------
49 ValueError
50 If dipole moment files are not suitable or mutually compatible.
51 """
52 atol_time = 1e-6
53 allclose_time = partial(np.allclose, rtol=0, atol=atol_time)
55 if isinstance(dipole_files, (Path, str)):
56 dipole_files = [dipole_files]
57 elif len(dipole_files) != 3:
58 raise ValueError('Three dipole files expected')
60 # Read files
61 data_i = []
62 Data = namedtuple('Data', ['kickv', 'time_t', 'dm_tv'])
63 for fpath in dipole_files:
64 time_t, dm_tv, kick_v = _read_and_clean_dipole_moment_data(fpath)
65 dm_tv -= dm_tv[0]
66 # Find kick direction
67 kickv = np.argmax(np.abs(kick_v))
68 for v in range(3):
69 if v != kickv and not np.isclose(kick_v[v], 0, atol=1e-24):
70 raise ValueError(f'Kick in {fpath} not aligned in '
71 f'Cartesian direction')
72 data_i.append(Data(kickv, time_t, dm_tv / kick_v[kickv]))
74 # Check file consistency
75 time_t = data_i[0].time_t
76 if not all([allclose_time(d.time_t, time_t) for d in data_i]):
77 raise ValueError('Inconsistent times in dipole files')
79 # Assume diagonal response if single dipole file given
80 if len(data_i) == 1:
81 data = data_i[0]
82 dm_t = data.dm_tv[:, data.kickv]
83 # Build dipole tensor
84 dm_tvv = np.eye(3)[np.newaxis, :, :] * dm_t[:, np.newaxis, np.newaxis]
85 else:
86 # Check that all kick directions are included
87 kickv_i = [d.kickv for d in data_i]
88 for v in range(3):
89 if v not in kickv_i:
90 raise ValueError(f'Kick direction {"xyz"[v]} missing')
92 # Build dipole tensor
93 dm_tvv = np.zeros((len(time_t), 3, 3))
94 for data in data_i:
95 dm_tvv[:, data.kickv, :] = data.dm_tv
97 # Read atoms
98 if atoms_file is not None:
99 atoms = ase.io.read(atoms_file)
100 assert isinstance(atoms, Atoms)
101 else:
102 atoms = None
104 # Initialize response object
105 response = TimeDomainResponse(time_t, dm_tvv, broadening=broadening,
106 name=name, units='au')
107 response._atoms = atoms
108 return response
111def _read_and_clean_dipole_moment_data(filepath: Path_str, atol_time: float = 1e-6) \
112 -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
113 """Reads and cleans a time-dependent dipole moment file from a real-time
114 time-dependent density functional theory calculation.
116 Parameters
117 ----------
118 filepath
119 File path.
120 atol_time
121 Absolute tolerance for comparing time values.
123 Returns
124 -------
125 time_t
126 Time values with time zero set to the kick time; shape {t}.
127 dm_tv
128 Dipole moment values, data epsilon before kick discarded; shape {t}x{3}.
129 kick_v
130 Kick strength; shape {3}.
132 Raises
133 ------
134 ValueError
135 If dipole moment file contains multiple kicks or
136 if dipole moment file has varying time step.
137 """
138 kick_i, time_t, _, dm_tv = _read_dipole_moment_file(filepath)
140 # Check kick
141 if len(kick_i) > 1:
142 raise ValueError(f'Multiple kicks in {filepath}')
143 kick = kick_i[0]
145 # Set time zero to the kick time and discard data before kick
146 time_t -= kick.time
147 flt_t = time_t > 0 - atol_time
148 time_t = time_t[flt_t]
149 dm_tv = dm_tv[flt_t]
151 # Check time step
152 dt_t = time_t[1:] - time_t[:-1]
153 dt = dt_t[0]
154 if not np.allclose(dt_t, dt, rtol=0, atol=atol_time):
155 raise ValueError('Varying time step')
157 return time_t, dm_tv, kick.strength_v
160def _read_dipole_moment_file(filepath: Path_str, remove_duplicates: bool = True) \
161 -> Tuple[List[Kick], np.ndarray, np.ndarray, np.ndarray]:
162 """Reads a time-dependent dipole moment file generated by a real-time
163 time-dependent density functional theory calculation.
165 Parameters
166 ----------
167 filepath
168 File path.
169 remove_duplicates
170 If true, data for overlapping time instances is removed;
171 the data read first is kept.
173 Returns
174 -------
175 kick_i
176 List of kicks.
177 time_t
178 Time values; shape {t}.
179 norm_t
180 Norm values; shape {t}.
181 dm_tv
182 Dipole moment values; shape {t}x{3}.
184 Raises
185 ------
186 ValueError
187 If kick lines cannot be read.
188 """
189 # Search and read kicks
190 kick_re = re.compile(r'Kick = \['
191 r'(?P<k0>[-+0-9\.e\ ]+), '
192 r'(?P<k1>[-+0-9\.e\ ]+), '
193 r'(?P<k2>[-+0-9\.e\ ]+)\]')
194 time_re = re.compile(r'Time = (?P<time>[-+0-9\.e\ ]+)')
195 kick_i = []
196 with open(filepath, 'r') as f:
197 for line in f:
198 if line.startswith('# Kick'):
199 # Kick strength
200 m = kick_re.search(line)
201 if m is None:
202 raise ValueError('Kick line unreadable')
203 strength_v = np.array([float(m.group(f'k{v}')) for v in range(3)])
204 # Kick time
205 m = time_re.search(line)
206 if m is None:
207 time = 0.0
208 else:
209 time = float(m.group('time'))
210 kick_i.append(Kick(time, strength_v))
212 # Read data
213 time_t, data_ti = _read_td_file(filepath, remove_duplicates)
214 norm_t = data_ti[:, 0]
215 dm_tv = data_ti[:, 1:]
216 return kick_i, time_t, norm_t, dm_tv
219def _read_td_file(filepath: Path_str, remove_duplicates: bool = True) \
220 -> Tuple[np.ndarray, np.ndarray]:
221 """Reads a time data series from file.
223 The data file is assumed to be a plain-text file with the first column
224 being the time values and the remaining columns the data values.
226 Parameters
227 ----------
228 filepath
229 File path.
230 remove_duplicates
231 If true, data for overlapping time instances is removed;
232 the data read first is kept.
234 Returns
235 -------
236 time_t
237 Time values; shape `{t}`.
238 data_ti
239 Data values; shape {t}x{the number of data columns}.
240 """
241 # Read data
242 data_tj = np.loadtxt(filepath)
243 time_t = data_tj[:, 0]
244 data_ti = data_tj[:, 1:]
246 # Remove duplicates due to abruptly stopped and restarted calculation
247 if remove_duplicates:
248 flt_t = np.ones_like(time_t, dtype=bool)
249 maxtime = time_t[0]
250 for t in range(1, len(time_t)):
251 if time_t[t] > maxtime:
252 maxtime = time_t[t]
253 else:
254 flt_t[t] = False
255 time_t = time_t[flt_t]
256 data_ti = data_ti[flt_t]
257 return time_t, data_ti