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

1from __future__ import annotations 

2 

3import pytest 

4import numpy as np 

5 

6from gpaw.mpi import world 

7 

8 

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. 

18 

19 The end result is a spectrum, so we compare also to the result of SpectrumCalculator. 

20 

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 

26 

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 

37 

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 

42 

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 ) 

51 

52 calc.calculate_and_write(tmp_path / 'dipole_spectrum.npz', include_tcm=True) 

53 

54 spec_calc = SpectrumCalculator.from_file(dmfile, 

55 perturbation=response.perturbation) 

56 

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 

62 

63 # Compare to SpectrumCalculator result 

64 archive = np.load(tmp_path / 'dipole_spectrum.npz') 

65 spec_archive = np.load(tmp_path / 'spectrum.npz') 

66 

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) 

72 

73 # Integrate broadened TCM 

74 test_dm_wv = np.trapz(np.trapz(archive['dm_wouv'], energy_u, axis=2), energy_o, axis=1) 

75 

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) 

79 

80 

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. 

90 

91 Compare against the dipole obtained by convolution of the data in the dipole moment file. 

92 

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 

101 

102 pulses = [create_pulse(pulsefreq, 5.0, 10.0) for pulsefreq in [1.1, 1.3]] 

103 

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 

114 

115 perturbation = create_perturbation(perturbation) 

116 

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] 

120 

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] 

134 

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)] 

142 

143 # Convert to units of as and eÅ 

144 time_t *= au_to_as 

145 dm_pt *= au_to_eA 

146 

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 

152 

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) 

159 

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] 

165 

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 

175 

176 calc.calculate_and_write(tmp_path / 'dipole.npz', only_one_pulse=False) 

177 

178 if world.rank != 0: 

179 return 

180 

181 # Compare to SpectrumCalculator result 

182 archive = np.load(tmp_path / 'dipole.npz') 

183 np.testing.assert_allclose(time_t, archive['time_t'] * 1e3) 

184 

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)