1import numpy as np
2from numpy import sqrt, pi, exp, abs
3from scipy.special import erf
4
5import _gpaw
6from gpaw import debug
7from gpaw.utilities.tools import coordinates
8from gpaw.utilities import is_contiguous
9
10
11def Y_L(L, x, y, z, r2):
12    if L == 0:
13        return 0.28209479177387814
14    elif L == 1:
15        return 0.48860251190291992 * y
16    elif L == 2:
17        return 0.48860251190291992 * z
18    elif L == 3:
19        return 0.48860251190291992 * x
20    elif L == 4:
21        return 1.0925484305920792 * x * y
22    elif L == 5:
23        return 1.0925484305920792 * y * z
24    elif L == 6:
25        return 0.31539156525252005 * (3 * z * z - r2)
26    elif L == 7:
27        return 1.0925484305920792 * x * z
28    elif L == 8:
29        return 0.54627421529603959 * (x * x - y * y)
30
31
32def gauss_L(a, L, x, y, z, r2, exp_ar2):
33    if L == 0:
34        return sqrt(a**3 * 4) / pi * exp_ar2
35    elif L == 1:
36        return sqrt(a**5 * 5.333333333333333) / pi * y * exp_ar2
37    elif L == 2:
38        return sqrt(a**5 * 5.333333333333333) / pi * z * exp_ar2
39    elif L == 3:
40        return sqrt(a**5 * 5.333333333333333) / pi * x * exp_ar2
41    elif L == 4:
42        return sqrt(a**7 * 4.2666666666666666) / pi * x * y * exp_ar2
43    elif L == 5:
44        return sqrt(a**7 * 4.2666666666666666) / pi * y * z * exp_ar2
45    elif L == 6:
46        return sqrt(
47            a**7 * 0.35555555555555557) / pi * (3 * z * z - r2) * exp_ar2
48    elif L == 7:
49        return sqrt(a**7 * 4.2666666666666666) / pi * x * z * exp_ar2
50    elif L == 8:
51        return sqrt(a**7 * 1.0666666666666667) / pi * (x * x - y * y) * exp_ar2
52
53
54def gausspot_L(a, L, x, y, z, r, r2, erf_sar, exp_ar2):
55    if L == 0:
56        return 2.0 * 1.7724538509055159 * erf_sar / r
57    elif L == 1:
58        return 1.1547005383792515 * (1.7724538509055159 * erf_sar -
59                                     2 * sqrt(a) * r * exp_ar2) / r / r2 * y
60    elif L == 2:
61        return 1.1547005383792515 * (1.7724538509055159 * erf_sar -
62                                     2 * sqrt(a) * r * exp_ar2) / r / r2 * z
63    elif L == 3:
64        return 1.1547005383792515 * (1.7724538509055159 * erf_sar -
65                                     2 * sqrt(a) * r * exp_ar2) / r / r2 * x
66    elif L == 4:
67        return 0.5163977794943222 * (
68            5.3173615527165481 * erf_sar -
69            (6 + 4 *
70             (sqrt(a) * r)**2) * sqrt(a) * r * exp_ar2) / r / r2**2 * x * y
71    elif L == 5:
72        return 0.5163977794943222 * (
73            5.3173615527165481 * erf_sar -
74            (6 + 4 *
75             (sqrt(a) * r)**2) * sqrt(a) * r * exp_ar2) / r / r2**2 * y * z
76    elif L == 6:
77        return 0.14907119849998599 * (5.3173615527165481 * erf_sar -
78                                      (6 + 4 *
79                                       (sqrt(a) * r)**2) * sqrt(a) * r *
80                                      exp_ar2) / r / r2**2 * (3 * z * z - r2)
81    elif L == 7:
82        return 0.5163977794943222 * (
83            5.3173615527165481 * erf_sar -
84            (6 + 4 *
85             (sqrt(a) * r)**2) * sqrt(a) * r * exp_ar2) / r / r2**2 * x * z
86    elif L == 8:
87        return 0.2581988897471611 * (5.3173615527165481 * erf_sar -
88                                     (6 + 4 * (sqrt(a) * r)**2) * sqrt(a) * r *
89                                     exp_ar2) / r / r2**2 * (x * x - y * y)
90
91
92# end of computer generated code
93
94
95class Gaussian:
96    r"""Class offering several utilities related to the generalized gaussians.
97
98    Generalized gaussians are defined by::
99
100                       _____                           2
101                      /  1       l!         l+3/2  -a r   l  m
102       g (x,y,z) =   / ----- --------- (4 a)      e      r  Y (x,y,z),
103        L          \/  4 pi  (2l + 1)!                       l
104
105    where a is the inverse width of the gaussian, and Y_l^m is a real
106    spherical harmonic.
107    The gaussians are centered in the middle of input grid-descriptor."""
108    def __init__(self, gd, a=19., center=None):
109        self.gd = gd
110        self.xyz, self.r2 = coordinates(gd, center)
111        self.r = np.sqrt(self.r2)
112        self.set_width(a, center)
113        self.exp_ar2 = exp(-self.a * self.r2)
114        self.erf_sar = erf(sqrt(self.a) * self.r)
115
116    def set_width(self, a, center):
117        """Set exponent of exp-function to -a on the boundary."""
118        if center is None:
119            self.a = 4 * a * (self.gd.icell_cv**2).sum(1).max()
120        else:
121            cell_center = self.gd.cell_cv.sum(1) / 2
122            r_min = (cell_center - np.abs(center - cell_center)).min()
123            self.a = a / r_min**2
124
125    def get_gauss(self, L):
126        a = self.a
127        x, y, z = tuple(self.xyz)
128        r2 = self.r2
129        exp_ar2 = self.exp_ar2
130        return gauss_L(a, L, x, y, z, r2, exp_ar2)
131
132    def get_gauss_pot(self, L):
133        a = self.a
134        x, y, z = tuple(self.xyz)
135        r2 = self.r2
136        r = self.r
137        erf_sar = self.erf_sar
138        exp_ar2 = self.exp_ar2
139        return gausspot_L(a, L, x, y, z, r, r2, erf_sar, exp_ar2)
140
141    def get_moment(self, n, L):
142        r2 = self.r2
143        x, y, z = tuple(self.xyz)
144        return self.gd.integrate(n * Y_L(L, x, y, z, r2))
145
146    def remove_moment(self, n, L, q=None):
147        # Determine multipole moment
148        if q is None:
149            q = self.get_moment(n, L)
150
151        # Don't do anything if moment is less than the tolerance
152        if abs(q) < 1e-7:
153            return 0.
154
155        # Remove moment from input density
156        n -= q * self.get_gauss(L)
157
158        # Return correction
159        return q * self.get_gauss_pot(L)
160
161
162def gaussian_wave(r_vG, r0_v, sigma, k_v=None, A=None, dtype=float,
163                  out_G=None):
164    r"""Generates function values for atom-centered Gaussian waves.
165
166    ::
167
168                         _ _
169        _            / -|r-r0|^2 \           _ _
170      f(r) = A * exp( ----------- ) * exp( i k.r )
171                     \ 2 sigma^2 /
172
173    If the parameter A is not specified, the Gaussian wave is normalized::
174
175                                                  oo
176           /    ____        \ -3/2               /       _  2  2
177      A = (    /    '        )        =>    4 pi | dr |f(r)|  r  = 1
178           \ \/  pi   sigma /                    /
179                                                   0
180
181    Parameters:
182
183    r_vG: ndarray
184        Set of coordinates defining the grid positions.
185    r0_v: ndarray
186        Set of coordinates defining the center of the Gaussian envelope.
187    sigma: float
188        Specifies the spatial width of the Gaussian envelope.
189    k_v: ndarray or None
190        Set of reciprocal lattice coordinates defining the wave vector.
191        An argument of None is interpreted as the gamma point i.e. k_v=0.
192    A: float, complex or None
193        Specifies the amplitude of the Gaussian wave. Normalizes if None.
194    dtype: type, defaults to float
195        Specifies the output data type. Only returns the real-part if float.
196    out_G: ndarray or None
197        Optional pre-allocated buffer to fill in values. Allocates if None.
198
199    """
200    if k_v is None:
201        k_v = np.zeros(r0_v.shape)
202
203    if A is None:
204        # 4*pi*int(exp(-r^2/(2*sigma^2))^2 * r^2, r=0...infinity)
205        # = sigma^3*pi^(3/2) = 1/A^2 -> A = (sqrt(Pi)*sigma)^(-3/2)
206        A = 1 / (sigma * np.pi**0.5)**1.5
207
208    if debug:
209        assert is_contiguous(r_vG, float)
210        assert is_contiguous(r0_v, float)
211        assert is_contiguous(k_v, float)
212        assert r_vG.ndim >= 2 and r_vG.shape[0] > 0
213        assert r0_v.ndim == 1 and r0_v.shape[0] > 0
214        assert k_v.ndim == 1 and k_v.shape[0] > 0
215        assert (r_vG.shape[0], ) == r0_v.shape == k_v.shape
216        assert sigma > 0
217
218    if out_G is None:
219        out_G = np.empty(r_vG.shape[1:], dtype=dtype)
220    elif debug:
221        assert is_contiguous(out_G)
222        assert out_G.shape == r_vG.shape[1:]
223
224    # slice_v2vG = [slice(None)] + [np.newaxis]*3
225    # gw = lambda r_vG, r0_v, sigma, k_v, A=1/(sigma*np.pi**0.5)**1.5: \
226    #    * np.exp(-np.sum((r_vG-r0_v[slice_v2vG])**2, axis=0)/(2*sigma**2)) \
227    #    * np.exp(1j*np.sum(np.r_vG*k_v[slice_v2vG], axis=0)) * A
228    _gpaw.utilities_gaussian_wave(A, r_vG, r0_v, sigma, k_v, out_G)
229    return out_G
230