1import numpy as np
2import numbers
3from scipy.optimize import minimize
4from scipy.integrate import simps
5from gpaw import GPAW, PW
6from ase.parallel import parprint
7import scipy.integrate as integrate
8from scipy.interpolate import InterpolatedUnivariateSpline
9from ase.units import Hartree as Ha
10
11
12class ElectrostaticCorrections():
13    """
14    Calculate the electrostatic corrections for charged defects.
15    """
16    def __init__(self, pristine, charged,
17                 q=None, sigma=None, r0=None, dimensionality='3d'):
18        if isinstance(pristine, str):
19            pristine = GPAW(pristine, txt=None, parallel={'domain': 1})
20        if isinstance(charged, str):
21            charged = GPAW(charged, txt=None)
22        calc = GPAW(mode=PW(500, force_complex_dtype=True),
23                    kpts={'size': (1, 1, 1),
24                          'gamma': True},
25                    parallel={'domain': 1},
26                    symmetry='off',
27                    txt=None)
28        atoms = pristine.atoms.copy()
29        calc.initialize(atoms)
30
31        self.pristine = pristine
32        self.charged = charged
33        self.calc = calc
34        self.q = q
35        self.sigma = sigma
36        self.dimensionality = dimensionality
37
38        self.pd = self.calc.wfs.pd
39        self.L = self.pd.gd.cell_cv[2, 2]
40        if r0 is not None:
41            self.r0 = r0
42        elif dimensionality == '2d':
43            self.r0 = np.array([0, 0, self.L / 2])
44        else:
45            self.r0 = np.array([0, 0, 0])
46        self.z0 = self.r0[2]
47
48        self.G_Gv = self.pd.get_reciprocal_vectors(q=0,
49                                                   add_q=False)  # G in Bohr^-1
50        self.G2_G = self.pd.G2_qG[0]  # |\vec{G}|^2 in Bohr^-2
51        self.rho_G = self.calculate_gaussian_density()
52        self.Omega = abs(np.linalg.det(self.calc.density.gd.cell_cv))
53        self.data = None
54        self.El = None
55
56        # For the 2D case, we assume that the dielectric profile epsilon(z)
57        # follows the electronic density of the pristine system n(z).
58        density = self.pristine.get_pseudo_density(gridrefinement=2)
59        density_1d = density.mean(axis=(0, 1))
60        coarse_z = find_z(self.pristine.density.gd.refine())
61        fine_z = find_z(self.pd.gd.refine())
62        transformer = InterpolatedUnivariateSpline(coarse_z, density_1d)
63        self.density_1d = np.array([transformer(x) for x in fine_z])
64
65        # We potentially have two z axes -- one on the original calculation,
66        # and one on the calculator we just set up.
67        self.z_g = fine_z
68        self.density_z = coarse_z
69
70        self.G_z = find_G_z(self.pd)
71        self.G_parallel = np.unique(self.G_Gv[:, :2], axis=0)
72        self.GG = np.outer(self.G_z, self.G_z)  # G * Gprime
73
74        # We need the G_z vectors on the finely sampled grid, to calculate
75        # epsilon(G - G'), which goes up to 2G_max in magnitude.
76        self.G_z_fine = (np.fft.fftfreq(len(fine_z)) *
77                         2 * np.pi / (self.L) * len(fine_z))
78        self.index_array = self.get_index_array()
79
80        self.epsilons_z = None
81        self.Eli = None
82        self.Elp = None
83        self.epsilon_GG = None
84
85    def calculate_gaussian_density(self):
86        # Fourier transformed gaussian:
87        return self.q * np.exp(-0.5 * self.G2_G * self.sigma ** 2)
88
89    def set_epsilons(self, epsilons, epsilon_bulk=1):
90        """Set the bulk dielectric constant of the system corresponding to the
91        atomic configuration of the pristine system. This is used to calculate
92        the screened coulomb potential of the gaussian model distribution that
93        we use. In 2D, the dielectric constant is poorly defined. However, when
94        we do DFT calculations with periodic systems, we are always working
95        with pseudo bulk-like structures, due to the periodic boundary
96        conditions. It is therefore possible to calculate some dielectric
97        constant. The screening of the system must come from where there is an
98        electronic density, and we thus set
99
100        epsilon(z) = k * n(z) + epsilon_bulk,
101
102        for some constant of proprotionality k. This is determined by the
103        constraint that the total screening int dz epsilon(z) gives the correct
104        value.
105
106        For 3d systems, the trick is to set k = 0. In that case, we have the
107        constant dielectric function that we require.
108
109        Parameters:
110          epsilons: A float, or a list of two floats. If a single
111            float is given, the screening in-plane and out-of-plane is assumed
112            to be the same. If two floats are given, the first represents the
113            in-plane response, and the second represents the out-of-plane
114            response.
115          epsilon_bulk: The value of the screening far away from the system,
116            given in the same format as epsilons. In most cases, this will be
117            1, corresponding to vacuum.
118
119        Returns:
120            epsilons_z: dictionary containnig the normalized dielectric
121              profiles in the in-plane and out-of-plane directions along the z
122              axis.
123        """
124
125        # Reset the state of the object; the corrections need to be
126        # recalculated
127        self.Elp = None
128        self.Eli = None
129        self.epsilon_GG = None
130
131        def normalize(value):
132            if isinstance(value, numbers.Number):
133                value = [value, value]
134            elif len(value) == 1:
135                value = [value[0]] * 2
136            return np.array(value)
137
138        epsilons = normalize(epsilons)
139        self.epsilons = epsilons
140
141        if self.dimensionality == '2d':
142            eb = normalize(epsilon_bulk)
143        elif self.dimensionality == '3d':
144            eb = normalize(epsilons)
145        self.eb = eb
146
147        z = self.z_g
148        L = self.L
149        density_1d = self.density_1d
150        N = simps(density_1d, z)
151        epsilons_z = np.zeros((2,) + np.shape(density_1d))
152        epsilons_z += density_1d
153
154        # In-plane
155        epsilons_z[0] = epsilons_z[0] * (epsilons[0] - eb[0]) / N * L + eb[0]
156
157        # Out-of-plane
158        def objective_function(k):
159            k = k[0]
160            integral = simps(1 / (k * density_1d + eb[1]), z) / L
161            return np.abs(integral - 1 / epsilons[1])
162
163        test = minimize(objective_function, [1], method='Nelder-Mead')
164        assert test.success, "Unable to normalize dielectric profile"
165        k = test.x[0]
166        epsilons_z[1] = epsilons_z[1] * k + eb[1]
167        self.epsilons_z = {'in-plane': epsilons_z[0],
168                           'out-of-plane': epsilons_z[1]}
169        self.calculate_epsilon_GG()
170
171    def get_index_array(self):
172        """
173        Calculate the indices that map between the G vectors on the fine grid,
174        and the matrix defined by M(G, Gprime) = G - Gprime.
175
176        We use this to generate the matrix epsilon_GG = epsilon(G - Gprime)
177        based on knowledge of epsilon(G) on the fine grid.
178        """
179        G, Gprime = np.meshgrid(self.G_z, self.G_z)
180        difference_GG = G - Gprime
181        index_array = np.zeros(np.shape(difference_GG))
182        index_array[:] = np.nan
183        for idx, val in enumerate(self.G_z_fine):
184            mask = np.isclose(difference_GG, val)
185            index_array[mask] = idx
186        if np.isnan(index_array).any():
187            print("Missing indices found in mapping between plane wave "
188                  "sets. Something is wrong with your G-vectors!")
189            print(np.isnan(index_array).all())
190            print(np.where(np.isnan(index_array)))
191            print(self.G_z_fine)
192            print(difference_GG)
193            assert False
194        return index_array.astype(int)
195
196    def calculate_epsilon_GG(self):
197        assert self.epsilons_z is not None, ("You provide the dielectric "
198                                             "constant first!")
199        N = len(self.density_1d)
200        epsilon_par_G = np.fft.fft(self.epsilons_z['in-plane']) / N
201        epsilon_perp_G = np.fft.fft(self.epsilons_z['out-of-plane']) / N
202        epsilon_par_GG = epsilon_par_G[self.index_array]
203        epsilon_perp_GG = epsilon_perp_G[self.index_array]
204
205        self.epsilon_GG = {'in-plane': epsilon_par_GG,
206                           'out-of-plane': epsilon_perp_GG}
207
208    def calculate_periodic_correction(self):
209        G_z = self.G_z
210        if self.Elp is not None:
211            return self.Elp
212        if self.epsilon_GG is None:
213            self.calculate_epsilon_GG()
214        Elp = 0.0
215
216        for vector in self.G_parallel:
217            G2 = (np.dot(vector, vector) + G_z * G_z)
218            rho_G = self.q * np.exp(- G2 * self.sigma ** 2 / 2, dtype=complex)
219            if self.dimensionality == '2d':
220                rho_G *= np.exp(1j * self.z0 * G_z)
221            A_GG = (self.GG * self.epsilon_GG['out-of-plane']
222                    + np.dot(vector, vector) * self.epsilon_GG['in-plane'])
223            if np.allclose(vector, 0):
224                A_GG[0, 0] = 1  # The d.c. potential is poorly defined
225            V_G = np.linalg.solve(A_GG, rho_G)
226            if np.allclose(vector, 0):
227                parprint('Skipping G^2=0 contribution to Elp')
228                V_G[0] = 0
229            Elp += (rho_G * V_G).sum()
230
231        Elp *= 2.0 * np.pi * Ha / self.Omega
232
233        self.Elp = Elp
234        return Elp
235
236    def calculate_isolated_correction(self):
237        if self.Eli is not None:
238            return self.Eli
239
240        L = self.L
241        G_z = self.G_z
242
243        G, Gprime = np.meshgrid(G_z, G_z)
244        Delta_GG = G - Gprime
245        phase = Delta_GG * self.z0
246        gaussian = (Gprime ** 2 + G ** 2) * self.sigma * self.sigma / 2
247
248        prefactor = np.exp(1j * phase - gaussian)
249
250        if self.epsilon_GG is None:
251            self.calculate_epsilon_GG()
252
253        dE_GG_par = (self.epsilon_GG['in-plane']
254                     - self.eb[0] * np.eye(len(self.G_z)))
255        dE_GG_perp = (self.epsilon_GG['out-of-plane']
256                      - self.eb[1] * np.eye(len(self.G_z)))
257
258        def integrand(k):
259            K_inv_G = ((self.eb[0] * k ** 2 + self.eb[1] * G_z ** 2) /
260                       (1 - np.exp(-k * L / 2) * np.cos(L * G_z / 2)))
261            K_inv_GG = np.diag(K_inv_G)
262            D_GG = L * (K_inv_GG + dE_GG_perp * self.GG + dE_GG_par * k ** 2)
263            return (k * np.exp(-k ** 2 * self.sigma ** 2) *
264                    (prefactor * np.linalg.inv(D_GG)).sum())
265
266        I = integrate.quad(integrand, 0, np.inf, limit=500)
267        Eli = self.q * self.q * I[0] * Ha
268        self.Eli = Eli
269        return Eli
270
271    def calculate_potential_alignment(self):
272        V_neutral = -np.mean(self.pristine.get_electrostatic_potential(),
273                             (0, 1))
274        V_charged = -np.mean(self.charged.get_electrostatic_potential(),
275                             (0, 1))
276
277        z = self.density_z
278        z_model, V_model = self.calculate_z_avg_model_potential()
279        V_model = InterpolatedUnivariateSpline(z_model[:-1], V_model[:-1])
280        V_model = V_model(z)
281        self.V_model = V_model
282        Delta_V = self.average(V_model - V_charged + V_neutral, z)
283        return Delta_V
284
285    def calculate_z_avg_model_potential(self):
286
287        vox3 = self.calc.density.gd.cell_cv[2, :] / len(self.z_g)
288
289        # The grid is arranged with z increasing fastest, then y
290        # then x (like a cube file)
291
292        G_z = self.G_z[1:]
293        rho_Gz = self.q * np.exp(-0.5 * G_z * G_z * self.sigma * self.sigma)
294
295        zs = []
296        Vs = []
297        if self.dimensionality == '2d':
298            phase = np.exp(1j * (self.G_z * self.z0))
299            A_GG = (self.GG * self.epsilon_GG['out-of-plane'])
300            A_GG[0, 0] = 1
301            V_G = np.linalg.solve(A_GG,
302                                  phase * np.array([0] + list(rho_Gz)))[1:]
303        elif self.dimensionality == '3d':
304            V_G = rho_Gz / self.eb[1] / G_z ** 2
305
306        for z in self.z_g:
307            phase_G = np.exp(1j * (G_z * z))
308            V = (np.sum(phase_G * V_G)
309                 * Ha * 4.0 * np.pi / (self.Omega))
310            Vs.append(V)
311
312        V = (np.sum(V_G) * Ha * 4.0 * np.pi / (self.Omega))
313        zs = list(self.z_g) + [vox3[2]]
314        Vs.append(V)
315        return np.array(zs), np.array(Vs)
316
317    def average(self, V, z):
318        N = len(V)
319        if self.dimensionality == '3d':
320            middle = N // 2
321            points = range(middle - N // 8, middle + N // 8 + 1)
322            restricted = V[points]
323        elif self.dimensionality == '2d':
324            points = list(range(0, N // 8)) + list(range(7 * N // 8, N))
325        restricted = V[points]
326        V_mean = np.mean(restricted)
327        return V_mean
328
329    def calculate_corrected_formation_energy(self):
330        E_0 = self.pristine.get_potential_energy()
331        E_X = self.charged.get_potential_energy()
332        Eli = self.calculate_isolated_correction().real
333        Elp = self.calculate_periodic_correction().real
334        Delta_V = self.calculate_potential_alignment()
335        return E_X - E_0 - (Elp - Eli) + Delta_V * self.q
336
337    def calculate_uncorrected_formation_energy(self):
338        E_0 = self.pristine.get_potential_energy()
339        E_X = self.charged.get_potential_energy()
340        return E_X - E_0
341
342    def collect_electrostatic_data(self):
343        V_neutral = -np.mean(self.pristine.get_electrostatic_potential(),
344                             (0, 1)),
345        V_charged = -np.mean(self.charged.get_electrostatic_potential(),
346                             (0, 1)),
347        data = {'epsilon': self.eb[0],
348                'z': self.density_z,
349                'V_0': V_neutral,
350                'V_X': V_charged,
351                'Elc': self.Elp - self.Eli,
352                'D_V_mean': self.calculate_potential_alignment(),
353                'V_model': self.V_model,
354                'D_V': (self.V_model
355                        - V_neutral
356                        + V_charged)}
357        self.data = data
358        return data
359
360
361def find_G_z(pd):
362    G_Gv = pd.get_reciprocal_vectors(q=0)
363    mask = (G_Gv[:, 0] == 0) & (G_Gv[:, 1] == 0)
364    G_z = G_Gv[mask][:, 2]  # G_z vectors in Bohr^{-1}
365    return G_z
366
367
368def find_z(gd):
369    r3_xyz = gd.get_grid_point_coordinates()
370    nrz = r3_xyz.shape[3]
371    return r3_xyz[2].flatten()[:nrz]
372