Coverage for tests/unittests/test_voronoi.py: 51%

1from __future__ import annotations 

2 

3import numpy as np 

4import pytest 

5 

6from gpaw.mpi import world 

7 

8from tests import get_permanent_test_file 

9 

10from rhodent.voronoi import VoronoiWeightCalculator, VoronoiLCAOWeightCalculator, VoronoiLCAOReader 

11 

12 

13def write_reference_data(ref_voronoi, gpw_fname): 

14 assert world.size == 1, 'Run me in serial mode' 

15 

16 voronoi = VoronoiWeightCalculator.from_gpw(atom_projections=[[0], [1]], 

17 gpw_file=gpw_fname) 

18 voronoi.calculate_and_write(str(ref_voronoi)) 

19 

20 

21@pytest.mark.bigdata 

22@pytest.mark.parametrize('test_system', ['2H2', 'Na8', 'Ag8']) 

23def test_weight_nn_against_reference(tmp_path, ref_voronoi, gpw_fname): 

24 reference = np.load(ref_voronoi) 

25 

26 atom_projections = reference['atom_projections'] 

27 voronoi = VoronoiWeightCalculator.from_gpw(atom_projections=atom_projections, gpw_file=gpw_fname) 

28 

29 # Normal calculation 

30 for i, weight_nn in enumerate(voronoi): 

31 if world.rank == 0: 

32 ref_nn = reference['weight_inn'][i] 

33 np.testing.assert_allclose(ref_nn, weight_nn, atol=1e-12) 

34 else: 

35 assert weight_nn is None 

36 

37 # Calculate LCAO and save, then read 

38 voronoi.voronoi_lcao.calculate_and_write(str(tmp_path / 'voronoi_lcao.ulm')) 

39 

40 voronoi = VoronoiWeightCalculator(VoronoiLCAOReader(str(tmp_path / 'voronoi_lcao.ulm'))) 

41 for i, weight_nn in enumerate(voronoi): 

42 if world.rank == 0: 

43 ref_nn = reference['weight_inn'][i] 

44 np.testing.assert_allclose(ref_nn, weight_nn, atol=1e-12) 

45 else: 

46 assert weight_nn is None 

47 

48 

49@pytest.mark.parametrize('domain', [2]) # The interesting edge case is this intermediate value 

50@pytest.mark.parametrize('test_system', ['2H2', 'Na8']) # Skip Ag8 to save time 

51def test_domain(tmp_path, gpw_fname, domain): 

52 from gpaw.mpi import world 

53 

54 if domain > world.size or world.size % domain != 0 or world.size == domain: 

55 pytest.skip('World size not compatible with domain parallelization') 

56 

57 kwargs = dict(gpw_file=gpw_fname) 

58 atom_projections = [[1], [0, 2]] 

59 voronoi = VoronoiLCAOWeightCalculator(atom_projections=atom_projections, **kwargs) 

60 

61 # Calculate with default domain (all ranks) 

62 ref_weight_iMM = list(voronoi) 

63 

64 # write grid to file 

65 voronoi.voronoi_grid.write(tmp_path / 'ref_grid.npy', voronoi.calc, voronoi.log) 

66 

67 # Calculate with smaller domain 

68 kwargs['domain'] = domain 

69 voronoi = VoronoiLCAOWeightCalculator(atom_projections=atom_projections, **kwargs) 

70 test_weight_iMM = list(voronoi) 

71 

72 # write grid to file 

73 voronoi.voronoi_grid.write(tmp_path / 'grid.npy', voronoi.calc, voronoi.log) 

74 

75 if world.rank == 0: 

76 np.testing.assert_allclose(np.load(tmp_path / 'ref_grid.npy'), np.load(tmp_path / 'grid.npy')) 

77 np.testing.assert_allclose(ref_weight_iMM, test_weight_iMM, atol=1e-12) 

78 else: 

79 assert all(weight_MM is None for weight_MM in ref_weight_iMM) 

80 assert all(weight_MM is None for weight_MM in test_weight_iMM) 

81 

82 # Calculate with default domain, reading from file 

83 kwargs['domain'] = -1 

84 kwargs['voronoi_grid'] = str(tmp_path / 'grid.npy') 

85 voronoi = VoronoiLCAOWeightCalculator(atom_projections=atom_projections, **kwargs) 

86 test_weight_iMM = list(voronoi) 

87 

88 if world.rank == 0: 

89 np.testing.assert_allclose(ref_weight_iMM, test_weight_iMM, atol=1e-12) 

90 else: 

91 assert all(weight_MM is None for weight_MM in ref_weight_iMM) 

92 assert all(weight_MM is None for weight_MM in test_weight_iMM) 

93 

94 

95if __name__ == '__main__': 

96 for test_system in ['2H2', 'Na8', 'Ag8']: 

97 write_reference_data(get_permanent_test_file(test_system, 'ref_voronoi'), 

98 get_permanent_test_file(test_system, 'gpw_fname'))