Coverage for tests/unittests/density_matrices/test_exact_convolver.py: 49%

1from __future__ import annotations 

2 

3import numpy as np 

4import pytest 

5 

6from gpaw.mpi import world, broadcast 

7from rhodent.density_matrices.time import ExactFourierTransformer 

8from rhodent.utils import create_pulse, two_communicators 

9from rhodent.perturbation import create_perturbation 

10 

11 

12@pytest.fixture 

13def write_frho(tmp_path, ksd_fname, mock_response): 

14 tmp_path = broadcast(tmp_path if world.rank == 0 else None) # Make sure that the temp path is the same 

15 

16 response = mock_response() 

17 

18 # Write 

19 frho_fmt = str(tmp_path / 'w{freq:05.2f}-{reim}.npy') 

20 response.write_in_frequency(frho_fmt, frequencies=np.arange(2, 5.1, 0.5)) 

21 

22 return frho_fmt 

23 

24 

25@pytest.mark.parametrize('calc_size', [2, 4]) 

26@pytest.mark.parametrize('test_system', ['Na8', 'Ag8']) 

27def test_read_write_time(ksd_fname, write_frho, calc_size): 

28 if calc_size > world.size or world.size % calc_size != 0 or world.size == calc_size: 

29 pytest.skip('World size not compatible with calc size') 

30 

31 time = 30e3 

32 pulse = create_perturbation(create_pulse(3.3)) 

33 frho_fmt = write_frho 

34 perturbation = {'name': 'deltakick', 'strength': 1e-5} 

35 

36 kwargs = dict(frho_fmt=frho_fmt, ksd=ksd_fname, perturbation=perturbation, 

37 real=True, imag=True, 

38 freq_spacing=0.5, freq_min=2, freq_max=5) 

39 

40 _, calc_comm = two_communicators(-1, calc_size) 

41 # Read with calc_size=1 

42 transformer = ExactFourierTransformer.from_file(comm=calc_comm, **kwargs) 

43 ref_buffer = transformer.convolve(time, pulse, derivative=0) 

44 if calc_comm.rank != 0: 

45 assert ref_buffer is None 

46 else: 

47 assert ref_buffer is not None 

48 

49 transformer = ExactFourierTransformer.from_file(comm=calc_comm, **kwargs) 

50 

51 test_buffer = transformer.convolve(time, pulse, derivative=0) 

52 

53 if calc_comm.rank != 0: 

54 assert test_buffer is None 

55 else: 

56 assert test_buffer is not None 

57 

58 if calc_comm.rank != 0: 

59 return 

60 

61 np.testing.assert_equal(ref_buffer.real, test_buffer.real) 

62 np.testing.assert_equal(ref_buffer.imag, test_buffer.imag)