Coverage for gpaw/test/new/test_single_precision.py: 74%

1import pytest 

2import numpy as np 

3import subprocess 

4import sys 

5 

6from ase.build import molecule 

7 

8from gpaw.new.ase_interface import GPAW 

9from gpaw.gpu import cupy_is_fake 

10 

11 

12@pytest.mark.serial 

13@pytest.mark.parametrize('dtype', 

14 ['np.complex128', 

15 'np.complex64', 

16 'np.float64', 

17 'np.float32']) 

18@pytest.mark.parametrize('gpu', [False, True]) 

19def test_single_precision(dtype, gpu): 

20 try: 

21 result = subprocess.run( 

22 'GPAW_NO_C_EXTENSION=1 GPAW_CPUPY=1 ' 

23 f'python {__file__} {dtype} {gpu}', 

24 shell=True, capture_output=True, 

25 text=True, check=True) 

26 except subprocess.CalledProcessError as e: 

27 print(e.output) 

28 print(e.stderr) 

29 raise e 

30 print(result.stdout) 

31 

32 

33@pytest.mark.serial 

34@pytest.mark.gpu 

35@pytest.mark.skipif(cupy_is_fake, reason='No cupy') 

36@pytest.mark.parametrize('dtype', 

37 [np.complex128, 

38 np.complex64, 

39 np.float64, 

40 np.float32]) 

41def test_single_precision_gpu(dtype): 

42 run_single_precision(dtype=dtype, gpu='True') 

43 

44 

45@pytest.mark.serial 

46@pytest.mark.gpu 

47@pytest.mark.skipif(cupy_is_fake, reason='No cupy') 

48@pytest.mark.parametrize('dtype', 

49 [np.complex128, 

50 np.complex64, 

51 np.float64, 

52 np.float32]) 

53def test_single_precision_rmmdiis_gpu(dtype): 

54 run_single_precision_rmmdiis(dtype=dtype) 

55 

56 

57def run_single_precision(dtype, gpu): 

58 atoms = molecule('H2O') 

59 atoms.center(vacuum=2.5) 

60 

61 gpu = gpu == 'True' 

62 

63 atoms.calc = GPAW(xc={'name': 'LDA'}, 

64 symmetry='off', 

65 random=True, 

66 convergence={'energy': 1e-5, 

67 'forces': 1e-3, 

68 'eigenstates': 1e-6}, 

69 mode={'name': 'pw', 

70 'ecut': 200.0, 

71 'dtype': dtype}, 

72 parallel={'gpu': gpu} 

73 ) 

74 

75 e_pot = atoms.get_potential_energy() 

76 expected_e = 9.595593485742606 

77 

78 assert atoms.calc.wfs.dtype == dtype 

79 

80 assert e_pot == pytest.approx(expected_e, rel=1e-3) 

81 

82 

83def run_single_precision_rmmdiis(dtype): 

84 atoms = molecule('H2O') 

85 atoms.center(vacuum=2.5) 

86 

87 atoms.calc = GPAW(xc={'name': 'LDA'}, 

88 symmetry='off', 

89 convergence={'energy': 1e-5, 

90 'forces': 1e-3}, 

91 mode={'name': 'pw', 

92 'ecut': 200.0, 

93 'dtype': dtype}, 

94 eigensolver={'name': 'rmm-diis'}, 

95 parallel={'gpu': True} 

96 ) 

97 

98 e_pot = atoms.get_potential_energy() 

99 expected_e = 9.595593485742606 

100 

101 assert atoms.calc.wfs.dtype == dtype 

102 

103 assert e_pot == pytest.approx(expected_e, rel=1e-3) 

104 

105 

106if __name__ == '__main__': 

107 dtypes = {'np.float32': np.float32, 

108 'np.float64': np.float64, 

109 'np.complex64': np.complex64, 

110 'np.complex128': np.complex128} 

111 run_single_precision(dtype=dtypes[sys.argv[1]], gpu=sys.argv[2])