1import warnings
2
3from ase.parallel import parprint
4from ase.units import Bohr
5import numpy as np
6from scipy.special import erf
7
8from gpaw.solvation.poisson import (WeightedFDPoissonSolver,
9                                    ADM12PoissonSolver,
10                                    PolarizationPoissonSolver)
11from gpaw.solvation.dielectric import Dielectric
12from gpaw.grid_descriptor import GridDescriptor
13from gpaw.utilities.gauss import Gaussian
14from gpaw.fd_operators import Gradient
15from gpaw.test import equal
16
17nn = 3
18accuracy = 2e-10
19
20
21def gradient(gd, x, nn):
22    out = gd.empty(3)
23    for i in (0, 1, 2):
24        Gradient(gd, i, 1.0, nn).apply(x, out[i])
25    return out
26
27
28def make_gd(h, box, pbc):
29    diag = np.array([box] * 3)
30    cell = np.diag(diag)
31    grid_shape = tuple((diag / h * 2).astype(int))
32    return GridDescriptor(grid_shape, cell / Bohr, pbc)
33
34
35class MockDielectric(Dielectric):
36    def __init__(self, epsinf, nn):
37        Dielectric.__init__(self, epsinf)
38        self.nn = nn
39
40    def update(self, cavity):
41        grad_eps_vg = gradient(self.gd, self.eps_gradeps[0] - self.epsinf,
42                               self.nn)
43        for i in (0, 1, 2):
44            self.eps_gradeps[1 + i][...] = grad_eps_vg[i]
45
46
47def test_solvation_poisson():
48    box = 12.
49    gd = make_gd(h=.4, box=box, pbc=False)
50
51    def solve(ps, eps, rho):
52        phi = gd.zeros()
53        dielectric = MockDielectric(epsinf=eps.max(), nn=nn)
54        dielectric.set_grid_descriptor(gd)
55        dielectric.allocate()
56        dielectric.eps_gradeps[0][...] = eps
57        dielectric.update(None)
58        with warnings.catch_warnings():
59            # ignore production code warning for alternative psolvers
60            warnings.simplefilter("ignore")
61            solver = ps(nn=nn, relax='J', eps=accuracy)
62        solver.set_dielectric(dielectric)
63        solver.set_grid_descriptor(gd)
64        solver.solve(phi, rho)
65        return phi
66
67    psolvers = (WeightedFDPoissonSolver,
68                ADM12PoissonSolver,
69                PolarizationPoissonSolver)
70
71    # test neutral system with constant permittivity
72    parprint('neutral, constant permittivity')
73    epsinf = 80.
74    eps = gd.zeros()
75    eps.fill(epsinf)
76    qs = (-1., 1.)
77    shifts = (-1., 1.)
78    rho = gd.zeros()
79    phi_expected = gd.zeros()
80    for q, shift in zip(qs, shifts):
81        gauss_norm = q / np.sqrt(4 * np.pi)
82        gauss = Gaussian(gd, center=(box / 2. + shift) * np.ones(3) / Bohr)
83        rho += gauss_norm * gauss.get_gauss(0)
84        phi_expected += gauss_norm * gauss.get_gauss_pot(0) / epsinf
85
86    for ps in psolvers:
87        phi = solve(ps, eps, rho)
88        parprint(ps, np.abs(phi - phi_expected).max())
89        equal(phi, phi_expected, 1e-3)
90
91    # test charged system with constant permittivity
92    parprint('charged, constant permittivity')
93    epsinf = 80.
94    eps = gd.zeros()
95    eps.fill(epsinf)
96    q = -2.
97    gauss_norm = q / np.sqrt(4 * np.pi)
98    gauss = Gaussian(gd, center=(box / 2. + 1.) * np.ones(3) / Bohr)
99    rho_gauss = gauss_norm * gauss.get_gauss(0)
100    phi_gauss = gauss_norm * gauss.get_gauss_pot(0)
101    phi_expected = phi_gauss / epsinf
102
103    for ps in psolvers:
104        phi = solve(ps, eps, rho_gauss)
105        parprint(ps, np.abs(phi - phi_expected).max())
106        equal(phi, phi_expected, 1e-3)
107
108    # test non-constant permittivity
109    msgs = ('neutral, non-constant permittivity',
110            'charged, non-constant permittivity')
111    qss = ((-1., 1.), (2., ))
112    shiftss = ((-.4, .4), (-1., ))
113    epsshifts = (.0, -1.)
114
115    for msg, qs, shifts, epsshift in zip(msgs, qss, shiftss, epsshifts):
116        parprint(msg)
117        epsinf = 80.
118        gauss = Gaussian(gd, center=(box / 2. + epsshift) * np.ones(3) / Bohr)
119        eps = gauss.get_gauss(0)
120        eps = epsinf - eps / eps.max() * (epsinf - 1.)
121
122        rho = gd.zeros()
123        phi_expected = gd.zeros()
124        grad_eps = gradient(gd, eps - epsinf, nn)
125
126        for q, shift in zip(qs, shifts):
127            gauss = Gaussian(gd, center=(box / 2. + shift) * np.ones(3) / Bohr)
128            phi_tmp = gauss.get_gauss_pot(0)
129            xyz = gauss.xyz
130            fac = 2. * np.sqrt(gauss.a) * np.exp(-gauss.a * gauss.r2)
131            fac /= np.sqrt(np.pi) * gauss.r2
132            fac -= erf(np.sqrt(gauss.a) * gauss.r) / (gauss.r2 * gauss.r)
133            fac *= 2.0 * 1.7724538509055159
134            grad_phi = fac * xyz
135            laplace_phi = -4. * np.pi * gauss.get_gauss(0)
136            rho_tmp = -1. / (4. * np.pi) * (
137                (grad_eps * grad_phi).sum(0) + eps * laplace_phi)
138            norm = gd.integrate(rho_tmp)
139            rho_tmp /= norm * q
140            phi_tmp /= norm * q
141            rho += rho_tmp
142            phi_expected += phi_tmp
143
144        # PolarizationPoissonSolver does not pass this test
145        for ps in psolvers[:-1]:
146            phi = solve(ps, eps, rho)
147            parprint(ps, np.abs(phi - phi_expected).max())
148            equal(phi, phi_expected, 1e-3)
149