Coverage for gpaw/test/lcaotddft/test_laser.py: 100%

1import numpy as np 

2import pytest 

3 

4from ase.units import Bohr, Hartree 

5from gpaw.external import ConstantElectricField 

6from gpaw.lcaotddft import LCAOTDDFT 

7from gpaw.lcaotddft.dipolemomentwriter import DipoleMomentWriter 

8from gpaw.lcaotddft.laser import create_laser, register_custom_laser, Laser 

9from gpaw.mpi import world 

10from gpaw.tddft.units import as_to_au 

11 

12# Settings 

13dt = 20.0 

14N1 = 5 

15N = 5 + N1 

16kick_v = np.ones(3) * 1e-5 

17 

18 

19@pytest.mark.rttddft 

20@pytest.mark.parametrize('pulse', [ 

21 {'name': 'GaussianPulse', 'strength': 1e-5, 'time0': 0, 'frequency': 8.6, 

22 'sigma': 0.5, 'sincos': 'sin'}, 

23 {'name': 'SincPulse', 'strength': 1e-6, 'time0': 0, 'cutoff_freq': 15, 

24 'relative_t0': False}]) 

25def test_laser(gpw_files, in_tmp_dir, pulse): 

26 # Time-propagation calculation 

27 td_calc = LCAOTDDFT(gpw_files['na2_tddft_dzp'], txt='td.out') 

28 DipoleMomentWriter(td_calc, 'dm.dat') 

29 td_calc.absorption_kick(kick_v) 

30 td_calc.propagate(dt, N) 

31 

32 # Pulse 

33 direction = kick_v 

34 ext = ConstantElectricField(Hartree / Bohr, direction) 

35 

36 # Time-propagation calculation with pulse 

37 pulse = create_laser(pulse) 

38 td_calc = LCAOTDDFT(gpw_files['na2_tddft_dzp'], 

39 td_potential={'ext': ext, 'laser': pulse}, 

40 txt='tdpulse.out') 

41 DipoleMomentWriter(td_calc, 'dmpulse.dat') 

42 td_calc.propagate(dt, N1) 

43 td_calc.write('td.gpw', mode='all') 

44 # Restart 

45 td_calc = LCAOTDDFT('td.gpw', txt='tdpulse2.out') 

46 DipoleMomentWriter(td_calc, 'dmpulse.dat') 

47 td_calc.propagate(dt, N - N1) 

48 

49 # Convoluted dipole moment 

50 world.barrier() 

51 time_t = np.arange(0, dt * (N + 0.1), dt) * as_to_au 

52 pulse_t = pulse.strength(time_t) 

53 np.savetxt('pulse.dat', np.stack((time_t, pulse_t)).T) 

54 dm_tv = np.delete(np.loadtxt('dm.dat')[:, 2:], 1, axis=0) 

55 dm_tv /= np.linalg.norm(kick_v) 

56 pulsedm_tv = np.delete(np.loadtxt('dmpulse.dat')[:, 2:], N1, axis=0) 

57 

58 tol = 5e-6 

59 for v in range(3): 

60 pulsedmconv_t = np.convolve( 

61 dm_tv[:, v], pulse_t)[:(N + 1)] * dt * as_to_au 

62 np.savetxt('dmpulseconv%d.dat' % v, pulsedmconv_t) 

63 assert pulsedm_tv[:, v] == pytest.approx(pulsedmconv_t, abs=tol) 

64 

65 

66@pytest.mark.rttddft 

67def test_custom(gpw_files, in_tmp_dir): 

68 gpw_fname = gpw_files['na2_tddft_dzp'] 

69 

70 class RandomPulse(Laser): 

71 def __init__(self, strength): 

72 self.rng = np.random.default_rng(42) 

73 self.dict = dict(name='RandomPulse', strength=strength) 

74 self.s0 = strength 

75 

76 def strength(self, t): 

77 return self.s0 * self.rng.uniform(size=np.shape(t)) 

78 

79 def todict(self): 

80 return self.dict 

81 

82 # We should be able to run and restart using custom pulses 

83 register_custom_laser('RandomPulse', RandomPulse) 

84 

85 # Pulse 

86 direction = kick_v 

87 ext = ConstantElectricField(Hartree / Bohr, direction) 

88 

89 # Time-propagation calculation with pulse 

90 pulse = RandomPulse(1e-5) 

91 td_calc = LCAOTDDFT(gpw_fname, td_potential={'ext': ext, 'laser': pulse}, 

92 txt='tdpulse.out') 

93 td_calc.propagate(dt, 1) 

94 

95 td_calc.write('td.gpw', mode='all') 

96 

97 # Restart 

98 td_calc = LCAOTDDFT('td.gpw', txt='tdpulse2.out') 

99 

100 restart_pulse = td_calc.td_hamiltonian.td_potential.laser_i[0] 

101 

102 assert isinstance(restart_pulse, RandomPulse) 

103 np.testing.assert_equal(restart_pulse.todict(), pulse.todict())