Coverage for gpaw/elph/filter.py: 17%

1import numpy as np 

2import numpy.fft as fft 

3import numpy.linalg as la 

4 

5from ase import Atoms 

6import ase.units as units 

7 

8from gpaw.typing import ArrayND 

9 

10 

11def fourier_filter(atoms: Atoms, supercell: tuple, V1t_xG: ArrayND, 

12 components: str, criteria=0) -> None: 

13 """Fourier filter atomic gradients of the effective potential. 

14 

15 ... or any other property defined the same way. 

16 

17 This method is not tested. 

18 

19 Parameters 

20 ---------- 

21 atoms: Atoms 

22 Atoms object, primitive cell 

23 supercell: tuple 

24 Size of supercell given by the number of repetitions (l, m, n) of 

25 the small unit cell in each direction. 

26 V1t_xG: ndarray 

27 Array representation of atomic gradients of the effective potential 

28 in the supercell grid. 

29 components: str 

30 Fourier components to filter out (``normal`` or ``umklapp``). 

31 """ 

32 assert components in ['normal', 'umklapp'] 

33 # Grid shape 

34 shape = V1t_xG.shape[-3:] 

35 

36 # Primitive unit cells in Bohr/Bohr^-1 

37 cell_cv = atoms.get_cell() / units.Bohr 

38 reci_vc = 2 * np.pi * la.inv(cell_cv) 

39 norm_c = np.sqrt(np.sum(reci_vc**2, axis=0)) 

40 # Periodic BC array 

41 pbc_c = np.array(atoms.get_pbc(), dtype=bool) 

42 

43 # Supercell atoms and cell 

44 atoms_N = atoms * supercell 

45 supercell_cv = atoms_N.get_cell() / units.Bohr 

46 

47 # q-grid in units of the grid spacing (FFT ordering) 

48 q_cG = np.indices(shape).reshape(3, -1) 

49 q_c = np.array(shape)[:, np.newaxis] 

50 q_cG += q_c // 2 

51 q_cG %= q_c 

52 q_cG -= q_c // 2 

53 

54 # Locate q-points inside the Brillouin zone 

55 if criteria == 0: 

56 # Works for all cases 

57 # Grid spacing in direction of reciprocal lattice vectors 

58 h_c = np.sqrt(np.sum((2 * np.pi * la.inv(supercell_cv))**2, 

59 axis=0)) 

60 # XXX Why does a "*=" operation on q_cG not work here ?? 

61 q1_cG = q_cG * h_c[:, np.newaxis] / (norm_c[:, np.newaxis] / 2) 

62 mask_G = np.ones(np.prod(shape), dtype=bool) 

63 for i, pbc in enumerate(pbc_c): 

64 if not pbc: 

65 continue 

66 mask_G &= (-1. < q1_cG[i]) & (q1_cG[i] <= 1.) 

67 else: 

68 # 2D hexagonal lattice 

69 # Projection of q points onto the periodic directions. Only in 

70 # these directions do normal and umklapp processees make sense. 

71 q_vG = np.dot(q_cG[pbc_c].T, 

72 2 * np.pi * la.inv(supercell_cv).T[pbc_c]).T.copy() 

73 # Parametrize the BZ boundary in terms of the angle theta 

74 theta_G = np.arctan2(q_vG[1], q_vG[0]) % (np.pi / 3) 

75 phi_G = np.pi / 6 - np.abs(theta_G) 

76 qmax_G = norm_c[0] / 2 / np.cos(phi_G) 

77 norm_G = np.sqrt(np.sum(q_vG**2, axis=0)) 

78 # Includes point on BZ boundary with +1e-2 

79 mask_G = (norm_G <= qmax_G + 1e-2) 

80 

81 if components == 'umklapp': 

82 mask_G = ~mask_G 

83 

84 # Reshape to grid shape 

85 mask_G.shape = shape 

86 

87 for V1t_G in V1t_xG: 

88 # Fourier transform atomic gradient 

89 V1tq_G = fft.fftn(V1t_G) 

90 # Zero normal/umklapp components 

91 V1tq_G[mask_G] = 0.0 

92 # Fourier transform back 

93 V1t_G[:] = fft.ifftn(V1tq_G).real