1# -*- coding: utf-8
2"""This module defines Coulomb and XC kernels for the response model.
3"""
4
5import numpy as np
6from ase.dft import monkhorst_pack
7
8
9def get_coulomb_kernel(pd, N_c, truncation=None, q_v=None, wstc=None):
10    """Factory function that calls the specified flavour
11    of the Coulomb interaction"""
12
13    qG_Gv = pd.get_reciprocal_vectors(add_q=True)
14    if q_v is not None:
15        assert pd.kd.gamma
16        qG_Gv += q_v
17
18    if truncation is None:
19        if pd.kd.gamma and q_v is None:
20            v_G = np.zeros(len(pd.G2_qG[0]))
21            v_G[0] = 4 * np.pi
22            v_G[1:] = 4 * np.pi / (qG_Gv[1:]**2).sum(axis=1)
23        else:
24            v_G = 4 * np.pi / (qG_Gv**2).sum(axis=1)
25
26    elif truncation == '2D':
27        v_G = calculate_2D_truncated_coulomb(pd, q_v=q_v, N_c=N_c)
28        if pd.kd.gamma and q_v is None:
29            v_G[0] = 0.0
30
31    elif truncation == '1D':
32        v_G = calculate_1D_truncated_coulomb(pd, q_v=q_v, N_c=N_c)
33        if pd.kd.gamma and q_v is None:
34            v_G[0] = 0.0
35
36    elif truncation == '0D':
37        v_G = calculate_0D_truncated_coulomb(pd, q_v=q_v)
38        if pd.kd.gamma and q_v is None:
39            v_G[0] = 0.0
40
41    elif truncation == 'wigner-seitz':
42        v_G = wstc.get_potential(pd, q_v=q_v)
43        if pd.kd.gamma and q_v is None:
44            v_G[0] = 0.0
45    else:
46        raise ValueError('Truncation scheme %s not implemented' % truncation)
47
48    return v_G.astype(complex)
49
50
51def calculate_2D_truncated_coulomb(pd, q_v=None, N_c=None):
52    """ Simple 2D truncation of Coulomb kernel PRB 73, 205119.
53    The non-periodic direction is determined from k-point grid.
54    """
55
56    qG_Gv = pd.get_reciprocal_vectors(add_q=True)
57    if pd.kd.gamma:
58        if q_v is not None:
59            qG_Gv += q_v
60        else:  # only to avoid warning. Later set to zero in factory function
61            qG_Gv[0] = [1., 1., 1.]
62
63    # The non-periodic direction is determined from k-point grid
64    Nn_c = np.where(N_c == 1)[0]
65    Np_c = np.where(N_c != 1)[0]
66    if len(Nn_c) != 1:
67        # The k-point grid does not fit with boundary conditions
68        Nn_c = [2]  # choose reduced cell vectors 0, 1
69        Np_c = [0, 1]  # choose reduced cell vector 2
70    # Truncation length is half of cell vector in non-periodic direction
71    R = pd.gd.cell_cv[Nn_c[0], Nn_c[0]] / 2.
72
73    qGp_G = ((qG_Gv[:, Np_c[0]])**2 + (qG_Gv[:, Np_c[1]]**2))**0.5
74    qGn_G = qG_Gv[:, Nn_c[0]]
75
76    v_G = 4 * np.pi / (qG_Gv**2).sum(axis=1)
77    if np.allclose(qGn_G[0], 0) or pd.kd.gamma:
78        """sin(qGn_G * R) = 0 when R = L/2 and q_n = 0.0"""
79        v_G *= 1.0 - np.exp(-qGp_G * R) * np.cos(qGn_G * R)
80    else:
81        """Normal component of q is not zero"""
82        a_G = qGn_G / qGp_G * np.sin(qGn_G * R) - np.cos(qGn_G * R)
83        v_G *= 1. + np.exp(-qGp_G * R) * a_G
84
85    return v_G.astype(complex)
86
87
88def calculate_1D_truncated_coulomb(pd, q_v=None, N_c=None):
89    """ Simple 1D truncation of Coulomb kernel PRB 73, 205119.
90    The periodic direction is determined from k-point grid.
91    """
92
93    from scipy.special import j1, k0, j0, k1
94
95    qG_Gv = pd.get_reciprocal_vectors(add_q=True)
96    if pd.kd.gamma:
97        if q_v is not None:
98            qG_Gv += q_v
99        else:
100            raise ValueError(
101                'Presently, calculations only work with a small q in the '
102                'normal direction')
103
104    # The periodic direction is determined from k-point grid
105    Nn_c = np.where(N_c == 1)[0]
106    Np_c = np.where(N_c != 1)[0]
107    if len(Nn_c) != 2:
108        # The k-point grid does not fit with boundary conditions
109        Nn_c = [0, 1]    # Choose reduced cell vectors 0, 1
110        Np_c = [2]       # Choose reduced cell vector 2
111    # The radius is determined from area of non-periodic part of cell
112    Acell_cv = pd.gd.cell_cv[Nn_c, :][:, Nn_c]
113    R = (np.linalg.det(Acell_cv) / np.pi)**0.5
114
115    qGnR_G = (qG_Gv[:, Nn_c[0]]**2 + qG_Gv[:, Nn_c[1]]**2)**0.5 * R
116    qGpR_G = abs(qG_Gv[:, Np_c[0]]) * R
117    v_G = 4 * np.pi / (qG_Gv**2).sum(axis=1)
118    v_G *= (1.0 + qGnR_G * j1(qGnR_G) * k0(qGpR_G)
119            - qGpR_G * j0(qGnR_G) * k1(qGpR_G))
120
121    return v_G.astype(complex)
122
123
124def calculate_0D_truncated_coulomb(pd, q_v=None):
125    """ Simple spherical truncation of the Coulomb interaction
126    PRB 73, 205119
127    """
128
129    qG_Gv = pd.get_reciprocal_vectors(add_q=True)
130    if pd.kd.gamma:
131        if q_v is not None:
132            qG_Gv += q_v
133        else:  # Only to avoid warning. Later set to zero in factory function
134            qG_Gv[0] = [1., 1., 1.]
135    # The radius is determined from volume of cell
136    R = (3 * np.linalg.det(pd.gd.cell_cv) / (4 * np.pi))**(1. / 3.)
137
138    qG2_G = (qG_Gv**2).sum(axis=1)
139
140    v_G = 4 * np.pi / qG2_G
141    v_G *= 1.0 - np.cos(qG2_G**0.5 * R)
142
143    return v_G
144
145
146def get_integrated_kernel(pd, N_c, truncation=None, N=100, reduced=False):
147    from scipy.special import j1, k0, j0, k1
148
149    B_cv = 2 * np.pi * pd.gd.icell_cv
150    Nf_c = np.array([N, N, N])
151    if reduced:
152        # Only integrate periodic directions if truncation is used
153        Nf_c[np.where(N_c == 1)[0]] = 1
154    q_qc = monkhorst_pack(Nf_c) / N_c
155    q_qc += pd.kd.ibzk_kc[0]
156    q_qv = np.dot(q_qc, B_cv)
157
158    if truncation is None:
159        V_q = 4 * np.pi / np.sum(q_qv**2, axis=1)
160    elif truncation == '2D':
161        # The non-periodic direction is determined from k-point grid
162        Nn_c = np.where(N_c == 1)[0]
163        Np_c = np.where(N_c != 1)[0]
164        if len(Nn_c) != 1:
165            # The k-point grid does not fit with boundary conditions
166            Nn_c = [2]  # choose reduced cell vectors 0, 1
167            Np_c = [0, 1]  # choose reduced cell vector 2
168        # Truncation length is half of cell vector in non-periodic direction
169        R = pd.gd.cell_cv[Nn_c[0], Nn_c[0]] / 2.
170
171        qp_q = ((q_qv[:, Np_c[0]])**2 + (q_qv[:, Np_c[1]]**2))**0.5
172        qn_q = q_qv[:, Nn_c[0]]
173
174        V_q = 4 * np.pi / (q_qv**2).sum(axis=1)
175        a_q = qn_q / qp_q * np.sin(qn_q * R) - np.cos(qn_q * R)
176        V_q *= 1. + np.exp(-qp_q * R) * a_q
177    elif truncation == '1D':
178        # The non-periodic direction is determined from k-point grid
179        Nn_c = np.where(N_c == 1)[0]
180        Np_c = np.where(N_c != 1)[0]
181
182        if len(Nn_c) != 2:
183            # The k-point grid does not fit with boundary conditions
184            Nn_c = [0, 1]    # Choose reduced cell vectors 0, 1
185            Np_c = [2]       # Choose reduced cell vector 2
186        # The radius is determined from area of non-periodic part of cell
187        Acell_cv = pd.gd.cell_cv[Nn_c, :][:, Nn_c]
188        R = (np.linalg.det(Acell_cv) / np.pi)**0.5
189
190        qnR_q = (q_qv[:, Nn_c[0]]**2 + q_qv[:, Nn_c[1]]**2)**0.5 * R
191        qpR_q = abs(q_qv[:, Np_c[0]]) * R
192        V_q = 4 * np.pi / (q_qv**2).sum(axis=1)
193        V_q *= (1.0 + qnR_q * j1(qnR_q) * k0(qpR_q)
194                - qpR_q * j0(qnR_q) * k1(qpR_q))
195    elif truncation == '0D' or 'wigner-seitz':
196        R = (3 * np.linalg.det(pd.gd.cell_cv) / (4 * np.pi))**(1. / 3.)
197        q2_q = (q_qv**2).sum(axis=1)
198        V_q = 4 * np.pi / q2_q
199        V_q *= 1.0 - np.cos(q2_q**0.5 * R)
200
201    return np.sum(V_q) / len(V_q), np.sum(V_q**0.5) / len(V_q)
202