Coverage for gpaw/__init__.py: 72%

1# Copyright (C) 2003 CAMP 

2# Please see the accompanying LICENSE file for further information. 

3"""Main gpaw module.""" 

4from __future__ import annotations 

5import os 

6import sys 

7import contextlib 

8from pathlib import Path 

9from typing import List, Any, TYPE_CHECKING 

10import warnings 

11 

12 

13__version__ = '25.1.1b1' 

14__ase_version_required__ = '3.23.0' 

15 

16__all__ = ['GPAW', 

17 'Mixer', 'MixerSum', 'MixerDif', 'MixerSum2', 

18 'MixerFull', 

19 'CG', 'Davidson', 'RMMDIIS', 'DirectLCAO', 

20 'PoissonSolver', 

21 'FermiDirac', 'MethfesselPaxton', 'MarzariVanderbilt', 

22 'PW', 'LCAO', 'FD', 

23 'restart'] 

24 

25boolean_envvars = { 

26 'GPAW_NEW', 

27 'GPAW_CPUPY', 

28 'GPAW_USE_GPUS', 

29 'GPAW_TRACE', 

30 'GPAW_NO_C_EXTENSION', 

31 'GPAW_MPI4PY'} 

32allowed_envvars = { 

33 *boolean_envvars, 

34 'GPAW_MPI_OPTIONS', 

35 'GPAW_MPI', 

36 'GPAW_SETUP_PATH'} 

37 

38is_gpaw_python = '_gpaw' in sys.builtin_module_names 

39dry_run = 0 

40 

41# When type-checking or running pytest, we want the debug-wrappers enabled: 

42debug: bool = (TYPE_CHECKING or 

43 'pytest' in sys.modules or 

44 bool(sys.flags.debug)) 

45 

46if debug: 

47 for var in os.environ: 

48 if var.startswith('GPAW') and var not in allowed_envvars: 

49 warnings.warn(f'Unknown GPAW environment varable: {var}') 

50 

51 

52@contextlib.contextmanager 

53def disable_dry_run(): 

54 """Context manager for temporarily disabling dry-run mode. 

55 

56 Useful for skipping exit in the GPAW constructor. 

57 """ 

58 global dry_run 

59 size = dry_run 

60 dry_run = 0 

61 yield 

62 dry_run = size 

63 

64 

65def get_scipy_version(): 

66 import scipy 

67 # This is in a function because we don't like to have the scipy 

68 # import at module level 

69 return [int(x) for x in scipy.__version__.split('.')[:2]] 

70 

71 

72if 'OMP_NUM_THREADS' not in os.environ: 

73 os.environ['OMP_NUM_THREADS'] = '1' 

74 

75 

76class ConvergenceError(Exception): 

77 pass 

78 

79 

80class KohnShamConvergenceError(ConvergenceError): 

81 pass 

82 

83 

84class PoissonConvergenceError(ConvergenceError): 

85 pass 

86 

87 

88class KPointError(Exception): 

89 pass 

90 

91 

92class BadParallelization(Exception): 

93 """Error indicating missing parallelization support.""" 

94 pass 

95 

96 

97def get_libraries() -> dict[str, str]: 

98 import gpaw.cgpaw as cgpaw 

99 libraries: dict[str, str] = {} 

100 if hasattr(cgpaw, 'lxcXCFunctional'): 

101 libraries['libxc'] = getattr(cgpaw, 'libxc_version', '2.x.y') 

102 else: 

103 libraries['libxc'] = '' 

104 return libraries 

105 

106 

107def parse_arguments(argv): 

108 from argparse import (ArgumentParser, REMAINDER, 

109 RawDescriptionHelpFormatter) 

110 # With gpaw-python BLAS symbols are in global scope and we need to 

111 # ensure that NumPy and SciPy use symbols from their own dependencies 

112 if is_gpaw_python: 

113 old_dlopen_flags = sys.getdlopenflags() 

114 sys.setdlopenflags(old_dlopen_flags | os.RTLD_DEEPBIND) 

115 

116 if is_gpaw_python: 

117 sys.setdlopenflags(old_dlopen_flags) 

118 

119 version = sys.version.replace('\n', '') 

120 p = ArgumentParser(usage='%(prog)s [OPTION ...] [-c | -m] SCRIPT' 

121 ' [ARG ...]', 

122 description='Run a parallel GPAW calculation.\n\n' 

123 f'Compiled with:\n Python {version}', 

124 formatter_class=RawDescriptionHelpFormatter) 

125 

126 p.add_argument('--command', '-c', action='store_true', 

127 help='execute Python string given as SCRIPT') 

128 p.add_argument('--module', '-m', action='store_true', 

129 help='run library module given as SCRIPT') 

130 p.add_argument('-W', metavar='argument', 

131 action='append', default=[], dest='warnings', 

132 help='warning control. See the documentation of -W for ' 

133 'the Python interpreter') 

134 p.add_argument('script', metavar='SCRIPT', 

135 help='calculation script') 

136 p.add_argument('options', metavar='ARG', 

137 help='arguments forwarded to SCRIPT', nargs=REMAINDER) 

138 

139 args = p.parse_args(argv[1:]) 

140 

141 if args.command and args.module: 

142 p.error('-c and -m are mutually exclusive') 

143 

144 sys.argv = [args.script] + args.options 

145 

146 for w in args.warnings: 

147 # Need to convert between python -W syntax to call 

148 # warnings.filterwarnings(): 

149 warn_args = w.split(':') 

150 assert len(warn_args) <= 5 

151 

152 if warn_args[0] == 'all': 

153 warn_args[0] = 'always' 

154 if len(warn_args) >= 3: 

155 # e.g. 'UserWarning' (string) -> UserWarning (class) 

156 warn_args[2] = globals().get(warn_args[2]) 

157 if len(warn_args) == 5: 

158 warn_args[4] = int(warn_args[4]) 

159 

160 warnings.filterwarnings(*warn_args, append=True) 

161 

162 return args 

163 

164 

165def __getattr__(attr: str) -> Any: 

166 for attr_getter in _lazy_import, _get_gpaw_env_vars: 

167 try: 

168 result = attr_getter(attr) 

169 except AttributeError: 

170 continue 

171 return globals().setdefault(attr, result) 

172 raise _module_attr_error(attr) 

173 

174 

175def __dir__() -> List[str]: 

176 """ 

177 Get the (1) normally-present module attributes, (2) lazily-imported 

178 objects, and (3) envrionmental variables starting with `GPAW_`. 

179 """ 

180 return list({*globals(), 

181 *all_lazy_imports, # From `_lazy_import()` 

182 *{*boolean_envvars, # From `_get_gpaw_env_vars()` 

183 *(var for var in os.environ if var.startswith('GPAW_'))}}) 

184 

185 

186def _module_attr_error(attr: str, *args, **kwargs) -> AttributeError: 

187 return AttributeError(f'{__getattr__.__module__}: ' 

188 f'no attribute named `.{attr}`', 

189 *args, **kwargs) 

190 

191 

192def _lazy_import(attr: str) -> Any: 

193 """ 

194 Implement the lazy importing of classes in submodules.""" 

195 import importlib 

196 

197 try: 

198 import_target = all_lazy_imports[attr] 

199 except KeyError: 

200 raise _module_attr_error(attr) from None 

201 

202 module, sep, target = import_target.rpartition('.') 

203 assert module and all(chunk.isidentifier() for chunk in module.split('.')) 

204 assert sep 

205 assert target.isidentifier() 

206 return getattr(importlib.import_module(module), target) 

207 

208 

209def _get_gpaw_env_vars(attr: str) -> bool | str: 

210 if attr in boolean_envvars: 

211 return bool(int(os.environ.get(attr) or 0)) 

212 if attr in allowed_envvars and attr in os.environ: 

213 return os.environ[attr] 

214 raise _module_attr_error(attr) 

215 

216 

217all_lazy_imports = dict( 

218 Mixer='gpaw.mixer.Mixer', 

219 MixerSum='gpaw.mixer.MixerSum', 

220 MixerDif='gpaw.mixer.MixerDif', 

221 MixerSum2='gpaw.mixer.MixerSum2', 

222 MixerFull='gpaw.mixer.MixerFull', 

223 

224 Davidson='gpaw.eigensolvers.Davidson', 

225 RMMDIIS='gpaw.eigensolvers.RMMDIIS', 

226 CG='gpaw.eigensolvers.CG', 

227 DirectLCAO='gpaw.eigensolvers.DirectLCAO', 

228 

229 PoissonSolver='gpaw.poisson.PoissonSolver', 

230 FermiDirac='gpaw.occupations.FermiDirac', 

231 MethfesselPaxton='gpaw.occupations.MethfesselPaxton', 

232 MarzariVanderbilt='gpaw.occupations.MarzariVanderbilt', 

233 FD='gpaw.wavefunctions.fd.FD', 

234 LCAO='gpaw.wavefunctions.lcao.LCAO', 

235 PW='gpaw.wavefunctions.pw.PW') 

236 

237 

238# Make sure e.g. GPAW_NEW=0 will set GPAW_NEW=False 

239# (`__getattr__()` magic handles the other boolean environment 

240# variables, but GPAW_NEW is used within the same script, so it needs to 

241# concretely exist in the namespace) 

242GPAW_NEW = _get_gpaw_env_vars('GPAW_NEW') 

243 

244if os.uname().machine == 'wasm32': 

245 GPAW_NO_C_EXTENSION = True 

246 

247 

248class BroadcastImports: 

249 def __enter__(self): 

250 from gpaw._broadcast_imports import broadcast_imports 

251 self._context = broadcast_imports 

252 return self._context.__enter__() 

253 

254 def __exit__(self, *args): 

255 self._context.__exit__(*args) 

256 

257 

258broadcast_imports = BroadcastImports() 

259 

260 

261def main(): 

262 with broadcast_imports: 

263 import runpy 

264 

265 # Apparently we need the scipy.linalg import for compatibility? 

266 import scipy.linalg # noqa: F401 

267 

268 for attr in all_lazy_imports: 

269 __getattr__(attr) 

270 

271 gpaw_args = parse_arguments(sys.argv) 

272 # The normal Python interpreter puts . in sys.path, so we also do that: 

273 sys.path.insert(0, '.') 

274 # Stacktraces can be shortened by running script with 

275 # PyExec_AnyFile and friends. Might be nicer 

276 if gpaw_args.command: 

277 d = {'__name__': '__main__'} 

278 exec(gpaw_args.script, d, d) 

279 elif gpaw_args.module: 

280 # Python has: python [-m MOD] [-c CMD] [SCRIPT] 

281 # We use a much better way: gpaw-python [-m | -c] SCRIPT 

282 runpy.run_module(gpaw_args.script, run_name='__main__') 

283 else: 

284 runpy.run_path(gpaw_args.script, run_name='__main__') 

285 

286 

287if debug: 

288 import numpy as np 

289 np.seterr(over='raise', divide='raise', invalid='raise', under='ignore') 

290 oldempty = np.empty 

291 oldempty_like = np.empty_like 

292 

293 def empty(*args, **kwargs): 

294 a = oldempty(*args, **kwargs) 

295 try: 

296 a.fill(np.nan) 

297 except ValueError: 

298 a.fill(42) 

299 return a 

300 

301 def empty_like(*args, **kwargs): 

302 a = oldempty_like(*args, **kwargs) 

303 try: 

304 a.fill(np.nan) 

305 except ValueError: 

306 a.fill(-42) 

307 return a 

308 

309 np.empty = empty # type: ignore[misc] 

310 np.empty_like = empty_like 

311 

312if TYPE_CHECKING: 

313 from gpaw.new.ase_interface import GPAW 

314elif GPAW_NEW: 

315 all_lazy_imports['GPAW'] = 'gpaw.new.ase_interface.GPAW' 

316else: 

317 all_lazy_imports['GPAW'] = 'gpaw.calculator.GPAW' 

318 

319all_lazy_imports['get_calculation_info'] = 'gpaw.calcinfo.get_calculation_info' 

320 

321 

322def restart(filename, Class=None, **kwargs): 

323 if Class is None: 

324 from gpaw import GPAW as Class 

325 calc = Class(filename, **kwargs) 

326 atoms = calc.get_atoms() 

327 return atoms, calc 

328 

329 

330def read_rc_file(): 

331 home = os.environ.get('HOME') 

332 if home is not None: 

333 rc = os.path.join(home, '.gpaw', 'rc.py') 

334 if os.path.isfile(rc): 

335 # Read file in ~/.gpaw/rc.py 

336 with open(rc) as fd: 

337 exec(fd.read()) 

338 

339 

340def initialize_data_paths(): 

341 try: 

342 setup_paths[:0] = os.environ['GPAW_SETUP_PATH'].split(os.pathsep) 

343 except KeyError: 

344 pass 

345 

346 

347def standard_setup_paths() -> list[str | Path]: 

348 try: 

349 import gpaw_data 

350 except ModuleNotFoundError: 

351 return [] 

352 else: 

353 return [gpaw_data.datapath()] 

354 

355 

356setup_paths = standard_setup_paths() 

357read_rc_file() 

358initialize_data_paths() 

359 

360 

361def RMM_DIIS(*args, **kwargs): 

362 import warnings 

363 from gpaw import RMMDIIS 

364 warnings.warn('Please use RMMDIIS instead of RMM_DIIS') 

365 return RMMDIIS(*args, **kwargs)