Coverage for gpaw/utilities/elpa.py: 94%

1import atexit 

2import numpy as np 

3import gpaw.cgpaw as cgpaw 

4 

5 

6def _elpaconstants(): 

7 consts = cgpaw.pyelpa_constants() 

8 assert consts[0] == 0, f'ELPA_OK is {consts[0]}, expected 0' 

9 return {'elpa_ok': consts[0], 

10 '1stage': consts[1], 

11 '2stage': consts[2]} 

12 

13 

14class LibElpa: 

15 @staticmethod 

16 def have_elpa(): 

17 return hasattr(cgpaw, 'pyelpa_allocate') 

18 

19 _elpa_initialized = False 

20 

21 @staticmethod 

22 def api_version(): 

23 return cgpaw.pyelpa_version() 

24 

25 @classmethod 

26 def ensure_elpa_initialized(cls): 

27 if not cls._elpa_initialized: 

28 cgpaw.pyelpa_init() 

29 atexit.register(cgpaw.pyelpa_uninit) 

30 cls._elpa_initialized = True 

31 

32 def __init__(self, desc, nev=None, solver='1stage'): 

33 if nev is None: 

34 nev = desc.gshape[0] 

35 

36 ptr = np.zeros(1, np.intp) 

37 

38 if not self.have_elpa(): 

39 raise ImportError('GPAW is not running in parallel or otherwise ' 

40 'not compiled with Elpa support') 

41 

42 if desc.nb != desc.mb: 

43 raise ValueError('Row and column block size must be ' 

44 'identical to support Elpa') 

45 

46 self.ensure_elpa_initialized() 

47 cgpaw.pyelpa_allocate(ptr) 

48 self._ptr = ptr 

49 cgpaw.pyelpa_set_comm(ptr, 

50 desc.blacsgrid.comm.get_c_object()) 

51 self._parameters = {} 

52 

53 self._consts = _elpaconstants() 

54 elpasolver = self._consts[solver] 

55 

56 bg = desc.blacsgrid 

57 self.elpa_set(na=desc.gshape[0], 

58 local_ncols=desc.shape[0], 

59 local_nrows=desc.shape[1], 

60 nblk=desc.mb, 

61 process_col=bg.myrow, # XXX interchanged 

62 process_row=bg.mycol, 

63 blacs_context=bg.context) 

64 self.elpa_set(nev=nev, solver=elpasolver) 

65 self.desc = desc 

66 

67 cgpaw.pyelpa_setup(self._ptr) 

68 

69 @property 

70 def description(self): 

71 solver = self._parameters['solver'] 

72 if solver == self._consts['1stage']: 

73 pretty = 'Elpa one-stage solver' 

74 else: 

75 assert solver == self._consts['2stage'] 

76 pretty = 'Elpa two-stage solver' 

77 return pretty 

78 

79 @property 

80 def nev(self): 

81 return self._parameters['nev'] 

82 

83 def _is_complex(self, array): 

84 if array.dtype == np.complex128: 

85 return True 

86 if array.dtype == np.float64: 

87 return False 

88 

89 raise TypeError(f'Unsupported dtype {array.dtype} for Elpa interface') 

90 

91 def diagonalize(self, A, C, eps): 

92 assert self.nev == len(eps) 

93 self.desc.checkassert(A) 

94 self.desc.checkassert(C) 

95 cgpaw.pyelpa_diagonalize(self._ptr, A, C, eps, self._is_complex(A)) 

96 

97 def general_diagonalize(self, A, S, C, eps, is_already_decomposed=0): 

98 assert self.nev == len(eps) 

99 self.desc.checkassert(A) 

100 self.desc.checkassert(S) 

101 self.desc.checkassert(C) 

102 cgpaw.pyelpa_general_diagonalize(self._ptr, A, S, C, eps, 

103 is_already_decomposed, 

104 self._is_complex(A)) 

105 

106 def elpa_set(self, **kwargs): 

107 for key, value in kwargs.items(): 

108 # print('pyelpa_set {}={}'.format(key, value)) 

109 cgpaw.pyelpa_set(self._ptr, key, value) 

110 self._parameters[key] = value 

111 

112 def __repr__(self): 

113 return f'LibElpa({self._parameters})' 

114 

115 def __del__(self): 

116 if hasattr(self, '_ptr'): 

117 # elpa_deallocate has no error flag so we don't check it 

118 cgpaw.pyelpa_deallocate(self._ptr) 

119 self._ptr[0] = 0