Coverage for gpaw/wannier/functions.py: 78%

1from ase import Atoms 

2 

3from gpaw.calculator import GPAW 

4from gpaw.typing import Array3D 

5 

6 

7class WannierFunctions: 

8 def __init__(self, 

9 atoms: Atoms, 

10 centers, 

11 value, 

12 U_knn, 

13 n1=0, 

14 spin=0): 

15 self.atoms = atoms 

16 self.centers = centers 

17 self.U_knn = U_knn 

18 self.value = value 

19 self.n1 = 0 

20 self.spin = spin 

21 

22 def centers_as_atoms(self): 

23 return self.atoms + Atoms(f'X{len(self.centers)}', self.centers) 

24 

25 def get_function(self, calc: GPAW, n: int) -> Array3D: 

26 wf = calc.wfs.gd.zeros() 

27 for m, u in enumerate(self.U_knn[0][:, n]): 

28 wf += u * calc.wfs.get_wave_function_array(n=self.n1 + m, 

29 s=self.spin, 

30 k=0) 

31 return wf