Coverage for gpaw/gaunt.py: 100%

1from __future__ import annotations 

2 

3from typing import Dict 

4 

5import numpy as np 

6 

7from gpaw.typing import Array3D, Array4D 

8 

9_gaunt: Dict[int, np.ndarray] = {} 

10_nabla: Dict[int, np.ndarray] = {} 

11_super_gaunt: Dict[int, np.ndarray] = {} 

12 

13 

14def gaunt(lmax: int = 2) -> Array3D: 

15 r"""Gaunt coefficients 

16 

17 ::: 

18 

19 ^ ^ -- L ^ 

20 Y (r) Y (r) = > G Y (r) 

21 L L -- L L L 

22 1 2 L 1 2 

23 """ 

24 

25 if lmax in _gaunt: 

26 return _gaunt[lmax] 

27 

28 Lmax = (lmax + 1)**2 

29 L2max = (2 * lmax + 1)**2 

30 

31 from gpaw.spherical_harmonics import YL, gam 

32 G_LLL = np.zeros((Lmax, L2max, L2max)) 

33 for L1 in range(Lmax): 

34 for L2 in range(L2max): 

35 for L in range(L2max): 

36 r = 0.0 

37 for c1, n1 in YL[L1]: 

38 for c2, n2 in YL[L2]: 

39 for c, n in YL[L]: 

40 nx = n1[0] + n2[0] + n[0] 

41 ny = n1[1] + n2[1] + n[1] 

42 nz = n1[2] + n2[2] + n[2] 

43 r += c * c1 * c2 * gam(nx, ny, nz) 

44 G_LLL[L1, L2, L] = r 

45 _gaunt[lmax] = G_LLL 

46 return G_LLL 

47 

48 

49def nabla(lmax: int = 2) -> Array3D: 

50 """Create the array of derivative integrals. 

51 

52 ::: 

53 

54 / ^ ^ 1-l' d l' ^ 

55 | dr Y (r) r --[r Y (r)] 

56 / L _ L' 

57 dr 

58 """ 

59 

60 if lmax in _nabla: 

61 return _nabla[lmax] 

62 

63 Lmax = (lmax + 1)**2 

64 from gpaw.spherical_harmonics import YL, gam 

65 Y_LLv = np.zeros((Lmax, Lmax, 3)) 

66 # Insert new values 

67 for L1 in range(Lmax): 

68 for L2 in range(Lmax): 

69 for v in range(3): 

70 r = 0.0 

71 for c1, n1 in YL[L1]: 

72 for c2, n2 in YL[L2]: 

73 n = [0, 0, 0] 

74 n[0] = n1[0] + n2[0] 

75 n[1] = n1[1] + n2[1] 

76 n[2] = n1[2] + n2[2] 

77 if n2[v] > 0: 

78 # apply derivative 

79 n[v] -= 1 

80 # add integral 

81 r += n2[v] * c1 * c2 * gam(n[0], n[1], n[2]) 

82 Y_LLv[L1, L2, v] = r 

83 _nabla[lmax] = Y_LLv 

84 return Y_LLv 

85 

86 

87def super_gaunt(lmax: int = 2) -> Array4D: 

88 r"""Product of two Gaunt coefficients. 

89 

90 This gives the coefficients to contractions of three spherical harmonics: 

91 __ 

92 ˰ ˰ ˰ \ L ˰ 

93 Y (r) Y (r) Y (r) = / G Y (r) 

94 L L L ‾‾ L L L L 

95 1 2 3 L 1 2 3 

96 

97 where: 

98 __ 

99 L \ L' L 

100 G = / G G 

101 L L L ‾‾ L L L' L 

102 1 2 3 L' 1 2 3 

103 """ 

104 if lmax in _super_gaunt: 

105 return _super_gaunt[lmax] 

106 

107 G1_LLL = gaunt(lmax) 

108 G2_LLL = gaunt(2 * lmax) 

109 

110 Lmax1 = G1_LLL.shape[0] # l=0 to l=lmax 

111 Lmax2 = G2_LLL.shape[0] # l=0 to l=2*lmax 

112 G_LLLL = np.einsum('ijk,klm->ijlm', 

113 G1_LLL[:, :Lmax1], 

114 G2_LLL[:, :Lmax2]) 

115 

116 _super_gaunt[lmax] = G_LLLL 

117 return G_LLLL