Coverage for gpaw/new/fd/hamiltonian.py: 60%

1import numpy as np 

2from gpaw.core import UGArray 

3from gpaw.core.arrays import DistributedArrays as XArray 

4from gpaw.fd_operators import Gradient, Laplace 

5from gpaw.new import zips 

6from gpaw.new.hamiltonian import Hamiltonian 

7 

8 

9class FDHamiltonian(Hamiltonian): 

10 def __init__(self, grid, *, kin_stencil=3, xp=np): 

11 self.grid = grid 

12 self._gd = grid._gd 

13 self.kin = Laplace(self._gd, -0.5, kin_stencil, grid.dtype, xp=xp) 

14 

15 # For MGGA: 

16 self.grad_v = [] 

17 

18 def apply_local_potential(self, 

19 vt_R: UGArray, 

20 psit_nR: XArray, 

21 out: XArray, 

22 ) -> None: 

23 assert isinstance(psit_nR, UGArray) 

24 assert isinstance(out, UGArray) 

25 self.kin(psit_nR, out) 

26 for p, o in zips(psit_nR.data, out.data): 

27 o += p * vt_R.data 

28 

29 def apply_mgga(self, 

30 dedtaut_R: UGArray, 

31 psit_nR: XArray, 

32 vt_nR: XArray) -> None: 

33 if len(self.grad_v) == 0: 

34 grid = psit_nR.desc 

35 self.grad_v = [ 

36 Gradient(grid._gd, v, n=3, dtype=grid.dtype) 

37 for v in range(3)] 

38 

39 tmp_R = psit_nR.desc.empty() 

40 for psit_R, out_R in zips(psit_nR, vt_nR): 

41 for grad in self.grad_v: 

42 grad(psit_R, tmp_R) 

43 tmp_R.data *= dedtaut_R.data 

44 grad(tmp_R, tmp_R) 

45 tmp_R.data *= 0.5 

46 out_R.data -= tmp_R.data 

47 

48 def create_preconditioner(self, blocksize, xp=np): 

49 from types import SimpleNamespace 

50 

51 from gpaw.preconditioner import Preconditioner as PC 

52 pc = PC(self._gd, self.kin, self.grid.dtype, blocksize, xp=xp) 

53 

54 def apply(psit, residuals, out, ekin_n=None): 

55 kpt = SimpleNamespace(phase_cd=psit.desc.phase_factor_cd) 

56 pc(residuals.data, kpt, out=out.data) 

57 

58 return apply 

59 

60 def calculate_kinetic_energy(self, wfs, skip_sum=False): 

61 e_kin = 0.0 

62 for f, psit_R in zip(wfs.myocc_n, wfs.psit_nX): 

63 if f > 1.0e-10: 

64 e_kin += f * psit_R.integrate(self.kin(psit_R), skip_sum).real 

65 if not skip_sum: 

66 e_kin = psit_R.desc.comm.sum_scalar(e_kin) 

67 e_kin = wfs.band_comm.sum_scalar(e_kin) 

68 return e_kin * wfs.spin_degeneracy