Coverage for tests/__init__.py: 96%

1from __future__ import annotations 

2 

3import os 

4import numpy as np 

5from pathlib import Path 

6 

7testdatapath = Path('tests/data') 

8 

9permanent_test_files: dict[str, dict[str, Path]] = { 

10 '2H2': { 

11 'gpw_fname': testdatapath / '2H2_unocc.gpw', # Large file; not in git repo 

12 'fdm_fname': testdatapath / '2H2_fdm.ulm', 

13 'wfs_fname': testdatapath / '2H2_wfs.ulm', 

14 'ksd_fname': testdatapath / '2H2_ksd.ulm', 

15 'wfssnap_fname': testdatapath / '2H2_wfssnap.ulm', 

16 'ref_density_matrix': testdatapath / '2H2_ref.npz', 

17 'ref_voronoi': testdatapath / '2H2_voronoi_ref.npz', 

18 }, 

19 'Na8': { 

20 'gpw_fname': testdatapath / 'Na8_unocc.gpw', # Large file; not in git repo 

21 'fdm_fname': testdatapath / 'Na8_fdm.ulm', # Large file; not in git repo 

22 'wfs_fname': testdatapath / 'Na8_wfs.ulm', # Large file; not in git repo 

23 'ksd_fname': testdatapath / 'Na8_ksd.ulm', 

24 'wfssnap_fname': testdatapath / 'Na8_wfssnap.ulm', 

25 'ref_density_matrix': testdatapath / 'Na8_ref.npz', 

26 'ref_voronoi': testdatapath / 'Na8_voronoi_ref.npz', 

27 'ref_hcdist': testdatapath / 'Na8_hcdist_ref.npz', 

28 'ref_energy': testdatapath / 'Na8_energy_ref.npz', 

29 }, 

30 'Ag8': { 

31 'gpw_fname': testdatapath / 'Ag8_unocc.gpw', # Large file; not in git repo 

32 'fdm_fname': testdatapath / 'Ag8_fdm.ulm', # Large file; not in git repo 

33 'wfs_fname': testdatapath / 'Ag8_wfs.ulm', # Large file; not in git repo 

34 'ksd_fname': testdatapath / 'Ag8_ksd.ulm', 

35 'wfssnap_fname': testdatapath / 'Ag8_wfssnap.ulm', 

36 'ref_density_matrix': testdatapath / 'Ag8_ref.npz', 

37 'ref_voronoi': testdatapath / 'Ag8_voronoi_ref.npz', 

38 'ref_hcdist': testdatapath / 'Ag8_hcdist_ref.npz', 

39 'ref_energy': testdatapath / 'Ag8_energy_ref.npz', 

40 'dm_sinc': testdatapath / 'Ag8_sinc_dm.dat', 

41 'dm_delta': testdatapath / 'Ag8_delta_dm.dat', 

42 }, 

43 'Na55': { 

44 # Large enough system to have partially occupied states 

45 'ksd_fname': testdatapath / 'Na55_ksd.ulm', 

46 }, 

47 'Ag201CO': { 

48 # Large enough system to have properly decaying plasmon 

49 'dm_gauss': testdatapath / 'Ag201CO_gauss_dm.dat', 

50 'dm_sinc': testdatapath / 'Ag201CO_sinc_dm.dat', 

51 'dm_delta': testdatapath / 'Ag201CO_delta_dm.dat', 

52 'ref_spectrum': testdatapath / 'Ag201CO_spectrum_ref.npz', 

53 }, 

54} 

55 

56 

57def wrap_test(source_fct): 

58 from inspect import signature 

59 

60 def copy(target_fct): 

61 # target_fct is just a dummy 

62 def calltest(*args, **kwargs): 

63 source_fct(*args, **kwargs) 

64 

65 calltest.__signature__ = signature(source_fct) 

66 return calltest 

67 

68 return copy 

69 

70 

71def get_permanent_test_file(test_system, key): 

72 assert test_system in permanent_test_files, f'{test_system} is an invalid test system' 

73 return permanent_test_files[test_system][key] 

74 

75 

76def frho(ksd_fname, fdm_fname, frho_dname): 

77 from gpaw.mpi import SerialCommunicator, world 

78 from gpaw.lcaotddft.frequencydensitymatrix import FrequencyDensityMatrixReader 

79 from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition 

80 from gpaw.tddft.units import au_to_eV 

81 

82 from rhodent.utils import add_fake_kpts 

83 

84 calc_comm = SerialCommunicator() 

85 loop_comm = world 

86 

87 # Output directory 

88 if world.rank == 0: 

89 if not os.path.isdir(frho_dname): 

90 os.makedirs(frho_dname) 

91 world.barrier() 

92 

93 # Load ksd and fdm 

94 ksd = KohnShamDecomposition(filename=str(ksd_fname)) 

95 ksd.distribute(calc_comm) 

96 add_fake_kpts(ksd) 

97 fdm = FrequencyDensityMatrixReader(str(fdm_fname), ksd.ksl, ksd.kpt_u) 

98 

99 ffreq_w = fdm.freq_w 

100 rw_i = [(r, w) for r in ['Re', 'Im'] for w in range(len(ffreq_w))] 

101 for i in range(loop_comm.rank, len(rw_i), loop_comm.size): 

102 reim, w = rw_i[i] 

103 ffreq = ffreq_w[w] 

104 freq = ffreq.freq * au_to_eV 

105 folding = ffreq.folding.folding 

106 if folding is None: 

107 fname = 'w%05.2f-%s.npy' % (freq, reim) 

108 else: 

109 width = ffreq.folding.width * au_to_eV 

110 fname = 'w%05.2f-%s-%.3f-%s.npy' % (freq, folding, width, reim) 

111 fpath = os.path.join(frho_dname, fname) 

112 if not os.path.exists(fpath): 

113 print('Calculate %s' % fpath) 

114 rho_uMM = fdm.read_FDrho(reim, [w])[0] 

115 rho_p = ksd.transform(rho_uMM)[0] 

116 np.save(fpath, rho_p)