def solve(self,
phi,
rho,
charge=None,
eps=None,
maxcharge=1e-6,
zero_initial_phi=False):
self._init()
if self.is_extended:
self.rho_g[:] = 0
if not self.extended['useprev']:
self.phi_g[:] = 0
self.timer.start('Extend array')
extend_array(self.gd_original, self.gd, rho, self.rho_g)
self.timer.stop('Extend array')
retval = self._solve(self.phi_g, self.rho_g, charge, eps,
maxcharge, zero_initial_phi)
self.timer.start('Deextend array')
deextend_array(self.gd_original, self.gd, phi, self.phi_g)
self.timer.stop('Deextend array')
return retval
else:
return self._solve(phi, rho, charge, eps, maxcharge,
zero_initial_phi)
def calculate_induced_field(self,
from_density='comp',
gridrefinement=2,
extend_N_cd=None,
deextend=False,
poissonsolver=None,
gradient_n=3):
if self.has_field and \
from_density == self.field_from_density and \
self.Frho_wg is not None:
Frho_wg = self.Frho_wg
gd = self.fieldgd
else:
Frho_wg, gd = self.get_induced_density(from_density,
gridrefinement)
# Always extend a bit to get field without jumps
if extend_N_cd is None:
extend_N = max(8, 2**int(np.ceil(np.log(gradient_n) / np.log(2.))))
extend_N_cd = extend_N * np.ones(shape=(3, 2), dtype=np.int)
deextend = True
# Extend grid
oldgd = gd
egd, cell_cv, move_c = \
extended_grid_descriptor(gd, extend_N_cd=extend_N_cd)
Frho_we = egd.zeros((self.nw, ), dtype=self.dtype)
for w in range(self.nw):
extend_array(gd, egd, Frho_wg[w], Frho_we[w])
Frho_wg = Frho_we
gd = egd
if not deextend:
# TODO: this will make atoms unusable with original grid
self.atoms.set_cell(cell_cv, scale_atoms=False)
move_atoms(self.atoms, move_c)
# Allocate arrays
Fphi_wg = gd.zeros((self.nw, ), dtype=self.dtype)
Fef_wvg = gd.zeros((
self.nw,
self.nv,
), dtype=self.dtype)
Ffe_wg = gd.zeros((self.nw, ), dtype=float)
# Poissonsolver
if poissonsolver is None:
poissonsolver = PoissonSolver(name='fd', eps=1e-20)
poissonsolver.set_grid_descriptor(gd)
for w in range(self.nw):
# TODO: better output of progress
# parprint('%d' % w)
calculate_field(gd,
Frho_wg[w],
self.Fbgef_v,
Fphi_wg[w],
Fef_wvg[w],
Ffe_wg[w],
poissonsolver=poissonsolver,
nv=self.nv,
gradient_n=gradient_n)
# De-extend grid
if deextend:
Frho_wo = oldgd.zeros((self.nw, ), dtype=self.dtype)
Fphi_wo = oldgd.zeros((self.nw, ), dtype=self.dtype)
Fef_wvo = oldgd.zeros((
self.nw,
self.nv,
), dtype=self.dtype)
Ffe_wo = oldgd.zeros((self.nw, ), dtype=float)
for w in range(self.nw):
deextend_array(oldgd, gd, Frho_wo[w], Frho_wg[w])
deextend_array(oldgd, gd, Fphi_wo[w], Fphi_wg[w])
deextend_array(oldgd, gd, Ffe_wo[w], Ffe_wg[w])
for v in range(self.nv):
deextend_array(oldgd, gd, Fef_wvo[w][v], Fef_wvg[w][v])
Frho_wg = Frho_wo
Fphi_wg = Fphi_wo
Fef_wvg = Fef_wvo
Ffe_wg = Ffe_wo
gd = oldgd
# Store results
self.has_field = True
self.field_from_density = from_density
self.fieldgd = gd
self.Frho_wg = Frho_wg
self.Fphi_wg = Fphi_wg
self.Fef_wvg = Fef_wvg
self.Ffe_wg = Ffe_wg
def solve(self, phi, rho, **kwargs):
self._init()
phi_small_fine_g = phi
rho_small_fine_g = rho.copy()
if self.use_coarse:
# 1.1. Coarse rho on the small grid
tmp_g = rho_small_fine_g
for coarser in self.coarser_i:
tmp_g = coarser.apply(tmp_g)
rho_small_coar_g = tmp_g
else:
rho_small_coar_g = rho_small_fine_g
# 1.2. Extend rho to the large grid
rho_large_coar_g = self.gd_large_coar.zeros()
extend_array(self.gd_small_coar, self.gd_large_coar,
rho_small_coar_g, rho_large_coar_g)
# 1.3 Solve potential on the large coarse grid
niter_large = self.ps_large_coar.solve(self.phi_large_coar_g,
rho_large_coar_g, **kwargs)
rho_large_coar_g = None
if not self.use_coarse:
deextend_array(self.gd_small_fine, self.gd_large_coar,
phi_small_fine_g, self.phi_large_coar_g)
return niter_large
if self.use_aux_grid:
# 2.1 De-extend the potential to the auxiliary grid
phi_aux_coar_g = self.gd_aux_coar.empty()
deextend_array(self.gd_aux_coar, self.gd_large_coar_from_aux,
phi_aux_coar_g, self.phi_large_coar_g)
else:
phi_aux_coar_g = self.phi_large_coar_g
# 3.1 Refine the potential
tmp_g = phi_aux_coar_g
for refiner in self.refiner_i:
tmp_g = refiner.apply(tmp_g)
phi_aux_coar_g = None
phi_aux_fine_g = tmp_g
# 3.2 Calculate the corresponding density with Laplace
# (the refined coarse density would not accurately match with
# the potential)
rho_aux_fine_g = self.gd_aux_fine.empty()
self.laplace_aux_fine.apply(phi_aux_fine_g, rho_aux_fine_g)
# 3.3 De-extend the potential and density to the small grid
cor_phi_small_fine_g = self.gd_small_fine.empty()
deextend_array(self.gd_small_fine, self.gd_aux_fine,
cor_phi_small_fine_g, phi_aux_fine_g)
phi_aux_fine_g = None
cor_rho_small_fine_g = self.gd_small_fine.empty()
deextend_array(self.gd_small_fine, self.gd_aux_fine,
cor_rho_small_fine_g, rho_aux_fine_g)
rho_aux_fine_g = None
# 3.4 Remove the correcting density and potential
rho_small_fine_g -= cor_rho_small_fine_g
phi_small_fine_g -= cor_phi_small_fine_g
# 3.5 Solve potential on the small grid
niter_small = self.ps_small_fine.solve(phi_small_fine_g,
rho_small_fine_g, **kwargs)
# 3.6 Correct potential and density
phi_small_fine_g += cor_phi_small_fine_g
# rho_small_fine_g += cor_rho_small_fine_g
return (niter_large, niter_small)