Coverage for eminus/config.py: 77.17%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Consolidated configuration module.""" 

4 

5import inspect 

6import numbers 

7import os 

8import pathlib 

9import sys 

10 

11from .logger import log 

12 

13 

14class ConfigClass: 

15 """Configuration class holding user specifiable variables. 

16 

17 An instance of this class will be set as the same name as this module. This will effectively 

18 make this module a singleton data class. 

19 """ 

20 

21 def __init__(self): 

22 """Initialize the ConfigClass object.""" 

23 self.backend = "torch" # Use Torch as backend if available, default to NumPy otherwise 

24 self.use_gpu = False # Disable GPU by default, since one may be restrictd by a small VRAM 

25 self.use_pylibxc = True # Use Libxc over PySCF if available since it is faster 

26 self.threads = None # Read threads from environment variables by default 

27 self.verbose = "INFO" # Only display warnings (and worse) by default 

28 self.load_env() # Load environment variables to overwrite the defaults 

29 

30 # ### Class properties ### 

31 

32 @property 

33 def backend(self): 

34 """Whether to use NumPY or a different backend if installed.""" 

35 return self._backend 

36 

37 @backend.setter 

38 def backend(self, value): 

39 self._backend = value.lower() 

40 if self._backend == "torch": 

41 try: 

42 from array_api_compat import torch 

43 except ImportError: 

44 self._backend = "numpy" 

45 else: 

46 torch.set_default_dtype(torch.float64) 

47 elif self._backend == "cupy": 

48 try: 

49 from array_api_compat import cupy # noqa: F401 

50 except ImportError: 

51 self._backend = "numpy" 

52 else: 

53 self._backend = "numpy" 

54 

55 @property 

56 def use_gpu(self): 

57 """Whether to use the GPU if available.""" 

58 return self._use_gpu 

59 

60 @use_gpu.setter 

61 def use_gpu(self, value): 

62 if self.backend == "torch": 

63 import torch 

64 

65 # When using set_default_device the whole runtime will use a device context manager 

66 if torch.cuda.is_available(): 

67 # Parse strings from environment variables 

68 if isinstance(value, str): 

69 try: 

70 value = int(value) 

71 except ValueError: 

72 value = value.lower() == "true" 

73 if value: 

74 torch.set_default_device("cuda") 

75 self._use_gpu = True 

76 else: 

77 torch.set_default_device("cpu") 

78 self._use_gpu = False 

79 else: 

80 self._use_gpu = False 

81 elif self.backend == "cupy": 

82 self._use_gpu = True 

83 else: 

84 self._use_gpu = False 

85 

86 @property 

87 def use_pylibxc(self): 

88 """Whether to use pylibxc or PySCF for functionals if both are installed.""" 

89 if self._use_pylibxc: 

90 try: 

91 import pylibxc # noqa: F401 

92 except ImportError: 

93 pass 

94 else: 

95 return True 

96 return False 

97 

98 @use_pylibxc.setter 

99 def use_pylibxc(self, value): 

100 # Parse strings from environment variables 

101 if isinstance(value, str): 

102 try: 

103 value = int(value) 

104 except ValueError: 

105 value = value.lower() == "true" 

106 self._use_pylibxc = value 

107 

108 @property 

109 def threads(self): 

110 """Number of threads used in array calculations.""" 

111 if self._threads is None: 

112 try: 

113 if self.backend == "torch": 

114 from array_api_compat import torch 

115 

116 return torch.get_num_threads() 

117 # Read the OMP threads for the default operators 

118 return int(os.environ["OMP_NUM_THREADS"]) 

119 except KeyError: 

120 return None 

121 return int(self._threads) 

122 

123 @threads.setter 

124 def threads(self, value): 

125 self._threads = value 

126 if isinstance(value, numbers.Integral): 

127 if self.backend == "torch": 

128 from array_api_compat import torch 

129 

130 return torch.set_num_threads(value) 

131 os.environ["OMP_NUM_THREADS"] = str(value) 

132 return None 

133 

134 @property 

135 def verbose(self): 

136 """Logger verbosity level.""" 

137 return log.verbose 

138 

139 @verbose.setter 

140 def verbose(self, value): 

141 # Parse strings from environment variables 

142 try: # noqa: SIM105 

143 value = int(value) 

144 except ValueError: 

145 pass 

146 # Logic in setter to run it on initialization 

147 log.verbose = value 

148 

149 # ### Class methods ### 

150 

151 def info(self): 

152 """Print configuration and performance information.""" 

153 sys.stdout.write("--- Configuration infos ---\n") 

154 sys.stdout.write(f"Global verbosity : {self.verbose}\n") 

155 # Only print if PySCF or pylibxc is installed 

156 if not self.use_pylibxc: 

157 try: 

158 import pyscf # noqa: F401 

159 

160 sys.stdout.write("Libxc backend : PySCF\n") 

161 except ImportError: 

162 sys.stdout.write("Libxc backend : None\n") 

163 else: 

164 sys.stdout.write("Libxc backend : pylibxc\n") 

165 

166 sys.stdout.write( 

167 "\n--- Performance infos ---\n" 

168 f"Array backend : {self.backend}\n" 

169 f"Array device : {'GPU' if self.use_gpu else 'CPU'}\n" 

170 ) 

171 # Do not print threading information when using GPU 

172 if self.use_gpu: 

173 return 

174 # Check threads 

175 if self.threads is None: 

176 sys.stdout.write( 

177 "Array threads : 1\n" 

178 "INFO: No OMP_NUM_THREADS environment variable was found.\nTo improve " 

179 'performance, add "export OMP_NUM_THREADS=n" to your ".bashrc".\nMake sure to ' 

180 'replace "n", typically with the number of cores your CPU.\nTemporarily, you can ' 

181 'set them in your Python environment with "eminus.config.threads=n".\n' 

182 ) 

183 else: 

184 sys.stdout.write(f"Array threads : {self.threads}\n") 

185 

186 def load_env(self): 

187 """Load configurations from environment variables.""" 

188 properties = [] 

189 for name, value in inspect.getmembers(type(self)): 

190 if isinstance(value, property): 

191 properties.append(name) 

192 for prop in properties: 

193 if f"EMINUS_{prop.upper()}" in os.environ: 

194 setattr(self, prop, os.environ[f"EMINUS_{prop.upper()}"]) 

195 

196 

197if ( 

198 "sphinx-build" not in pathlib.Path(sys.argv[0]).name 

199 and "stubtest" not in pathlib.Path(sys.argv[0]).name 

200): 

201 # Do not initialize the class when Sphinx or stubtest is running 

202 # Since we set the class instance to the module name Sphinx would only document 

203 # the main docstring of the class without the properties 

204 sys.modules[__name__] = ConfigClass() 

205else: # pragma: no cover 

206 # Add mock variables for all properties and methods of the ConfigClass to the module 

207 # This allows IDEs to see that the module has said attribute 

208 # This also allows for stubtesting and documentation of these variables and functions 

209 backend = "" #: Whether to use NumPy or a different backend if installed. 

210 use_gpu = False #: Whether to use the GPU if available. 

211 use_pylibxc = False #: Whether to use pylibxc or PySCF for functionals if both are installed. 

212 threads = 0 #: Number of threads used in array calculations. 

213 verbose = "" #: Logger verbosity level. 

214 

215 def info(): 

216 """Print configuration and performance information.""" 

217 

218 def load_env(): 

219 """Load configurations from environment variables."""