Coverage for gpaw/new/c.py: 50%

1from typing import TYPE_CHECKING 

2 

3import gpaw.cgpaw as cgpaw 

4from gpaw import GPAW_NO_C_EXTENSION 

5from gpaw.new.timer import trace 

6 

7__all__ = ['GPU_AWARE_MPI'] 

8 

9GPU_AWARE_MPI = getattr(cgpaw, 'gpu_aware_mpi', False) 

10GPU_ENABLED = getattr(cgpaw, 'GPU_ENABLED', False) 

11 

12if not TYPE_CHECKING and not GPAW_NO_C_EXTENSION: 

13 from gpaw.cgpaw import (add_to_density, pw_insert, pw_precond, # noqa 

14 pwlfc_expand, symmetrize_ft) 

15 

16 if GPU_ENABLED: 

17 from gpaw.cgpaw import add_to_density_gpu # noqa 

18 from gpaw.cgpaw import (calculate_residuals_gpu, # noqa 

19 dH_aii_times_P_ani_gpu, evaluate_lda_gpu, 

20 evaluate_pbe_gpu, pw_amend_insert_realwf_gpu, 

21 pw_insert_gpu, pwlfc_expand_gpu, 

22 pw_norm_kinetic_gpu, pw_norm_gpu) 

23 

24 w = trace(gpu=True) 

25 add_to_density_gpu = w(add_to_density_gpu) 

26 calculate_residuals_gpu = w(calculate_residuals_gpu) 

27 dH_aii_times_P_ani_gpu = w(dH_aii_times_P_ani_gpu) 

28 evaluate_lda_gpu = w(evaluate_lda_gpu) 

29 evaluate_pbe_gpu = w(evaluate_pbe_gpu) 

30 pw_amend_insert_realwf_gpu = w(pw_amend_insert_realwf_gpu) 

31 pw_insert_gpu = w(pw_insert_gpu) 

32 pwlfc_expand_gpu = w(pwlfc_expand_gpu) 

33 pw_norm_kinetic_gpu = w(pw_norm_kinetic_gpu) 

34 pw_norm_gpu = w(pw_norm_gpu) 

35 else: 

36 from gpaw.purepython import (add_to_density_gpu, 

37 calculate_residuals_gpu, # noqa 

38 dH_aii_times_P_ani_gpu, evaluate_lda_gpu, 

39 evaluate_pbe_gpu, 

40 pw_amend_insert_realwf_gpu, 

41 pw_insert_gpu, pwlfc_expand_gpu, 

42 pw_norm_kinetic_gpu, pw_norm_gpu) 

43else: 

44 from gpaw.purepython import (add_to_density, pw_insert, pw_precond, # noqa 

45 pwlfc_expand, symmetrize_ft) 

46 from gpaw.purepython import (add_to_density_gpu, # noqa 

47 calculate_residuals_gpu, 

48 dH_aii_times_P_ani_gpu, evaluate_lda_gpu, 

49 evaluate_pbe_gpu, 

50 pw_amend_insert_realwf_gpu, 

51 pw_insert_gpu, pwlfc_expand_gpu, 

52 pw_norm_kinetic_gpu, pw_norm_gpu)