Coverage for tests/integration/test_dipole.py: 98%
97 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3import pytest
4import numpy as np
6from gpaw.mpi import world
9@pytest.mark.bigdata
10@pytest.mark.parametrize(['from_wfs', 'frequency_broadening'],
11 [(True, 0), (True, 0.1), (False, 0)])
12@pytest.mark.parametrize('test_system', ['Ag8'])
13def test_dipole_calculator_frequency_domain(tmp_path, from_wfs, frequency_broadening,
14 ksd_fname, wfs_fname, frho_dname,
15 dm_sinc, dm_delta):
16 """ Test that we can use the DipoleCalculator to compute
17 the response in the frequency domain from wave functions and from frho file.
19 The end result is a spectrum, so we compare also to the result of SpectrumCalculator.
21 This test also check that the integral of the TCM matches the spectrum.
22 """
23 from rhodent.response import ResponseFromWaveFunctions, ResponseFromFourierTransform
24 from rhodent.spectrum import SpectrumCalculator
25 from rhodent.calculators import DipoleCalculator
27 if from_wfs:
28 perturbation = {'name': 'SincPulse', 'strength': 1e-6,
29 'cutoff_freq': 4, 'time0': 5, 'relative_t0': True}
30 response = ResponseFromWaveFunctions(wfs_fname, ksd_fname, perturbation=perturbation)
31 dmfile = dm_sinc
32 else:
33 frho_fmt = str(frho_dname / 'w{freq:05.2f}-{reim}.npy')
34 perturbation = {'name': 'deltakick', 'strength': 1e-5}
35 response = ResponseFromFourierTransform(frho_fmt, ksd_fname, perturbation=perturbation)
36 dmfile = dm_delta
38 # Large, but rather coarse grid to integrate
39 energy_o = np.arange(-6, 2.001, 0.025)
40 energy_u = np.arange(-2, 4.001, 0.025)
41 sigma = 0.4
43 calc = DipoleCalculator(response=response,
44 frequencies=np.arange(0.05, 2.001, 0.05),
45 frequency_broadening=frequency_broadening,
46 voronoi=None,
47 energies_occ=energy_o,
48 energies_unocc=energy_u,
49 sigma=sigma,
50 )
52 calc.calculate_and_write(tmp_path / 'dipole_spectrum.npz', include_tcm=True)
54 spec_calc = SpectrumCalculator.from_file(dmfile,
55 perturbation=response.perturbation)
57 spec_calc.calculate_spectrum_and_write(tmp_path / 'spectrum.npz',
58 calc.frequencies,
59 frequency_broadening=frequency_broadening)
60 if world.rank != 0:
61 return
63 # Compare to SpectrumCalculator result
64 archive = np.load(tmp_path / 'dipole_spectrum.npz')
65 spec_archive = np.load(tmp_path / 'spectrum.npz')
67 atol = 0.1 * spec_archive['osc_wv'].max()
68 rtol = 0.01
69 np.testing.assert_allclose(archive['dm_wv'].imag * archive['osc_prefactor_w'][:, None],
70 spec_archive['osc_wv'],
71 rtol=rtol, atol=atol)
73 # Integrate broadened TCM
74 test_dm_wv = np.trapz(np.trapz(archive['dm_wouv'], energy_u, axis=2), energy_o, axis=1)
76 atol = 0.01 * np.abs(archive['dm_wv']).max()
77 rtol = 0.05
78 np.testing.assert_allclose(archive['dm_wv'], test_dm_wv, rtol=rtol, atol=atol)
81@pytest.mark.bigdata
82@pytest.mark.parametrize(['from_wfs', 'upscaling'],
83 [(True, 1), (True, 2), (True, 10), (True, 20), (False, 1)])
84@pytest.mark.parametrize('test_system', ['Ag8'])
85def test_dipole_calculator_time_domain(tmp_path, from_wfs, upscaling,
86 ksd_fname, wfs_fname, frho_dname,
87 dm_sinc, dm_delta):
88 """ Test that we can use the DipoleCalculator to compute
89 the response in the time domain from wave functions and from frho file.
91 Compare against the dipole obtained by convolution of the data in the dipole moment file.
93 Also test that we can compute from wave functions with upscaling and, that it is consistent
94 with the dipole moment file.
95 """
96 from rhodent.response import ResponseFromWaveFunctions, ResponseFromFourierTransform
97 from rhodent.calculators import DipoleCalculator
98 from rhodent.perturbation import create_perturbation
99 from rhodent.utils import create_pulse
100 from gpaw.tddft.units import au_to_as, au_to_eA
102 pulses = [create_pulse(pulsefreq, 5.0, 10.0) for pulsefreq in [1.1, 1.3]]
104 if from_wfs:
105 perturbation = {'name': 'SincPulse', 'strength': 1e-6,
106 'cutoff_freq': 4, 'time0': 5, 'relative_t0': True}
107 response = ResponseFromWaveFunctions(wfs_fname, ksd_fname, perturbation=perturbation)
108 dmfile = dm_sinc
109 else:
110 frho_fmt = str(frho_dname / 'w{freq:05.2f}-{reim}.npy')
111 perturbation = {'name': 'deltakick', 'strength': 1e-5}
112 response = ResponseFromFourierTransform(frho_fmt, ksd_fname, perturbation=perturbation)
113 dmfile = dm_delta
115 perturbation = create_perturbation(perturbation)
117 # Read dipole moment file and subtract static dipole
118 time_t, _, _, _, dm_t = np.loadtxt(dmfile, unpack=True)
119 dm_t -= dm_t[0]
121 # Remove any non-increasing times (due to restart of kick)
122 flt_t = np.ones_like(time_t, dtype=bool)
123 maxtime = time_t[0]
124 for t in range(1, len(time_t)):
125 if time_t[t] > maxtime:
126 maxtime = time_t[t]
127 else:
128 flt_t[t] = False
129 time_t = time_t[flt_t]
130 dm_t = dm_t[flt_t]
131 ndup = len(flt_t) - flt_t.sum()
132 print(f'Removed {ndup} duplicates')
133 dt = time_t[1] - time_t[0]
135 # Perform pulse convolution of the dipole moment file
136 padnt = 8096
137 pulse_pt = [pulse.strength(time_t) for pulse in pulses]
138 pulse_pw = np.fft.rfft(pulse_pt, axis=-1, n=padnt)
139 dm_w = perturbation.normalize_frequency_response(dm_t, time_t, padnt)
140 dm_pw = dm_w[None, :] * pulse_pw
141 dm_pt = np.fft.irfft(dm_pw, n=padnt, axis=-1)[:, :len(time_t)]
143 # Convert to units of as and eÅ
144 time_t *= au_to_as
145 dm_pt *= au_to_eA
147 # The dipole file contains all time steps, but the wave functions file contains
148 # the first time step, and then skips by 20
149 # This will guarantee that we can evaluate the dipole moment file convolution
150 # and DipoleCalculator result at the same times
151 times = np.arange(1, len(time_t), 20 / upscaling) * dt * au_to_as
153 # Extract the same times from the dipole moment file
154 # Only works if upscaling evenly divides 20
155 assert 20 % upscaling == 0
156 time_t = time_t[1::20//upscaling]
157 dm_pt = dm_pt[:, 1::20//upscaling]
158 np.testing.assert_allclose(time_t, times)
160 # And now, extract just a few times towards the end
161 flt_t = (time_t > 25e3) & (time_t < 27e3)
162 times = times[flt_t]
163 time_t = time_t[flt_t]
164 dm_pt = dm_pt[:, flt_t]
166 # Calculate using DipoleCalculator
167 calc = DipoleCalculator(response=response,
168 times=times,
169 pulses=pulses,
170 voronoi=None,
171 )
172 np.testing.assert_allclose(time_t, calc.times)
173 if from_wfs:
174 assert calc.density_matrices.rho_nn_conv.upscaling == upscaling
176 calc.calculate_and_write(tmp_path / 'dipole.npz', only_one_pulse=False)
178 if world.rank != 0:
179 return
181 # Compare to SpectrumCalculator result
182 archive = np.load(tmp_path / 'dipole.npz')
183 np.testing.assert_allclose(time_t, archive['time_t'] * 1e3)
185 # The error seems to be rather a phase shift (of less than one time step)
186 # so only check with absolute tolerance
187 atol = 1e-3 * np.abs(dm_pt).max()
188 np.testing.assert_allclose(dm_pt, archive['dm_ptv'][..., 2], rtol=0, atol=atol)