1# -*- coding: utf-8 -*- 2# Copyright (C) 2003 CAMP 3# Please see the accompanying LICENSE file for further information. 4 5"""This module defines a density class.""" 6 7from math import pi, sqrt 8 9import numpy as np 10from ase.units import Bohr 11 12from gpaw import debug 13from gpaw.mixer import get_mixer_from_keywords, MixerWrapper 14from gpaw.transformers import Transformer 15from gpaw.lfc import LFC, BasisFunctions 16from gpaw.wavefunctions.lcao import LCAOWaveFunctions 17from gpaw.utilities import (unpack2, unpack_atomic_matrices, 18 pack_atomic_matrices) 19from gpaw.utilities.partition import AtomPartition 20from gpaw.utilities.timing import nulltimer 21from gpaw.arraydict import ArrayDict 22 23 24class CompensationChargeExpansionCoefficients: 25 def __init__(self, setups, nspins): 26 self.setups = setups 27 self.nspins = nspins 28 29 def calculate(self, D_asp): 30 """Calculate multipole moments of compensation charges. 31 32 Returns the total compensation charge in units of electron 33 charge, so the number will be negative because of the 34 dominating contribution from the nuclear charge.""" 35 atom_partition = D_asp.partition 36 shape_a = [(setup.Delta_pL.shape[1],) for setup in self.setups] 37 Q_aL = atom_partition.arraydict(shape_a, dtype=float) 38 for a, D_sp in D_asp.items(): 39 setup = self.setups[a] 40 Q_L = np.dot(D_sp[:self.nspins].sum(0), setup.Delta_pL) 41 Q_L[0] += setup.Delta0 42 Q_aL[a] = Q_L 43 return Q_aL 44 45 def get_charge(self, Q_aL): 46 local_charge = sqrt(4 * pi) * sum(Q_L[0] for Q_L in Q_aL.values()) 47 return Q_aL.partition.comm.sum(local_charge) 48 49 50class NullBackgroundCharge: 51 charge = 0.0 52 53 def set_grid_descriptor(self, gd): 54 pass 55 56 def add_charge_to(self, rhot_g): 57 pass 58 59 def add_fourier_space_charge_to(self, pd, rhot_q): 60 pass 61 62 63class Density: 64 """Density object. 65 66 Attributes: 67 =============== ===================================================== 68 ``gd`` Grid descriptor for coarse grids. 69 ``finegd`` Grid descriptor for fine grids. 70 ``interpolate`` Function for interpolating the electron density. 71 ``mixer`` ``DensityMixer`` object. 72 =============== ===================================================== 73 74 Soft and smooth pseudo functions on uniform 3D grids: 75 ========== ========================================= 76 ``nt_sG`` Electron density on the coarse grid. 77 ``nt_sg`` Electron density on the fine grid. 78 ``nt_g`` Electron density on the fine grid. 79 ``rhot_g`` Charge density on the fine grid. 80 ``nct_G`` Core electron-density on the coarse grid. 81 ========== ========================================= 82 """ 83 84 def __init__(self, gd, finegd, nspins, collinear, charge, redistributor, 85 background_charge=None): 86 """Create the Density object.""" 87 88 self.gd = gd 89 self.finegd = finegd 90 self.nspins = nspins 91 self.collinear = collinear 92 self.charge = float(charge) 93 self.redistributor = redistributor 94 self.atomdist = None 95 96 self.ncomponents = self.nspins if self.collinear else 1 + 3 97 98 # This can contain e.g. a Jellium background charge 99 if background_charge is None: 100 background_charge = NullBackgroundCharge() 101 background_charge.set_grid_descriptor(self.finegd) 102 self.background_charge = background_charge 103 104 self.charge_eps = 1e-7 105 106 self.D_asp = None 107 self.Q = CompensationChargeExpansionCoefficients([], self.nspins) 108 self.Q_aL = None 109 110 self.nct_G = None 111 self.nt_xG = None 112 self.rhot_g = None 113 self.nt_xg = None 114 self.nt_sg = None 115 self.nt_vg = None 116 self.nt_g = None 117 118 self.atom_partition = None 119 120 self.setups = None 121 self.hund = None 122 self.magmom_av = None 123 124 self.fixed = False 125 # XXX at least one test will fail because None has no 'reset()' 126 # So we need DummyMixer I guess 127 self.mixer = None 128 self.set_mixer(None) 129 130 self.timer = nulltimer 131 self.error = None 132 self.nct = None 133 self.ghat = None 134 self.log = None 135 136 @property 137 def nt_sG(self): 138 return None if self.nt_xG is None else self.nt_xG[:self.nspins] 139 140 @property 141 def nt_vG(self): 142 return None if self.nt_xG is None else self.nt_xG[self.nspins:] 143 144 def __str__(self): 145 s = 'Densities:\n' 146 s += ' Coarse grid: {0}*{1}*{2} grid\n'.format(*self.gd.N_c) 147 s += ' Fine grid: {0}*{1}*{2} grid\n'.format(*self.finegd.N_c) 148 s += ' Total Charge: {0:.6f}'.format(self.charge) 149 if self.fixed: 150 s += '\n Fixed' 151 return s 152 153 def summary(self, atoms, magmom, log): 154 if self.nspins == 1: 155 return 156 try: 157 # XXX This doesn't always work, HGH, SIC, ... 158 sc = self.get_spin_contamination(atoms, int(magmom < 0)) 159 log('Spin contamination: %f electrons' % sc) 160 except (TypeError, AttributeError): 161 pass 162 163 def initialize(self, setups, timer, magmom_av, hund): 164 self.timer = timer 165 self.setups = setups 166 self.Q = CompensationChargeExpansionCoefficients(setups, self.nspins) 167 self.hund = hund 168 self.magmom_av = magmom_av 169 170 def reset(self): 171 # TODO: reset other parameters? 172 self.nt_xG = None 173 174 def set_positions_without_ruining_everything(self, spos_ac, 175 atom_partition): 176 rank_a = atom_partition.rank_a 177 # If both old and new atomic ranks are present, start a blank dict if 178 # it previously didn't exist but it will needed for the new atoms. 179 if (self.atom_partition is not None and 180 self.D_asp is None and (rank_a == self.gd.comm.rank).any()): 181 self.update_atomic_density_matrices( 182 self.setups.empty_atomic_matrix(self.ncomponents, 183 self.atom_partition)) 184 185 if (self.atom_partition is not None and self.D_asp is not None and 186 self.gd.comm.size > 1): 187 self.timer.start('Redistribute') 188 self.D_asp.redistribute(atom_partition) 189 self.timer.stop('Redistribute') 190 191 self.atom_partition = atom_partition 192 self.atomdist = self.redistributor.get_atom_distributions(spos_ac) 193 194 def set_positions(self, spos_ac, atom_partition): 195 self.set_positions_without_ruining_everything(spos_ac, atom_partition) 196 self.nct.set_positions(spos_ac, atom_partition) 197 self.ghat.set_positions(spos_ac, atom_partition) 198 self.mixer.reset() 199 200 self.nt_xg = None 201 self.nt_sg = None 202 self.nt_vg = None 203 self.nt_g = None 204 self.rhot_g = None 205 206 def calculate_pseudo_density(self, wfs): 207 """Calculate nt_sG from scratch. 208 209 nt_sG will be equal to nct_G plus the contribution from 210 wfs.add_to_density(). 211 """ 212 wfs.calculate_density_contribution(self.nt_xG) 213 self.nt_sG[:] += self.nct_G 214 215 def update_atomic_density_matrices(self, value): 216 if isinstance(value, dict): 217 tmp = self.setups.empty_atomic_matrix(self.ncomponents, 218 self.atom_partition) 219 tmp.update(value) 220 value = tmp 221 assert isinstance(value, ArrayDict) or value is None, type(value) 222 if value is not None: 223 value.check_consistency() 224 self.D_asp = value 225 226 def update(self, wfs): 227 self.timer.start('Density') 228 with self.timer('Pseudo density'): 229 self.calculate_pseudo_density(wfs) 230 with self.timer('Atomic density matrices'): 231 wfs.calculate_atomic_density_matrices(self.D_asp) 232 with self.timer('Multipole moments'): 233 comp_charge, _Q_aL = self.calculate_multipole_moments() 234 235 if isinstance(wfs, LCAOWaveFunctions): 236 self.timer.start('Normalize') 237 self.normalize(comp_charge) 238 self.timer.stop('Normalize') 239 240 self.timer.start('Mix') 241 self.mix(comp_charge) 242 self.timer.stop('Mix') 243 self.timer.stop('Density') 244 245 def normalize(self, comp_charge): 246 """Normalize pseudo density.""" 247 248 pseudo_charge = self.gd.integrate(self.nt_sG).sum() 249 250 if (pseudo_charge + self.charge + comp_charge - 251 self.background_charge.charge != 0): 252 if pseudo_charge != 0: 253 x = (self.background_charge.charge - self.charge - 254 comp_charge) / pseudo_charge 255 self.nt_xG *= x 256 else: 257 # Use homogeneous background. 258 # 259 # We have to use the volume per actual grid point, 260 # which is not the same as gd.volume as the latter 261 # includes ghost points. 262 volume = self.gd.get_size_of_global_array().prod() * self.gd.dv 263 total_charge = (self.charge + comp_charge - 264 self.background_charge.charge) 265 self.nt_sG[:] = -total_charge / volume 266 267 def mix(self, comp_charge): 268 assert isinstance(self.mixer, MixerWrapper), self.mixer 269 self.error = self.mixer.mix(self.nt_xG, self.D_asp) 270 assert self.error is not None, self.mixer 271 272 comp_charge = None 273 self.interpolate_pseudo_density(comp_charge) 274 self.calculate_pseudo_charge() 275 276 def calculate_multipole_moments(self): 277 D_asp = self.atomdist.to_aux(self.D_asp) 278 Q_aL = self.Q.calculate(D_asp) 279 self.Q_aL = Q_aL 280 return self.Q.get_charge(Q_aL), Q_aL 281 282 def get_initial_occupations(self, a): 283 # distribute charge on all atoms 284 # XXX interaction with background charge may be finicky 285 c = (self.charge - self.background_charge.charge) / len(self.setups) 286 M_v = self.magmom_av[a] 287 M = np.linalg.norm(M_v) 288 f_si = self.setups[a].calculate_initial_occupation_numbers( 289 M, self.hund, charge=c, 290 nspins=self.nspins if self.collinear else 2) 291 292 if self.collinear: 293 if M_v[2] < 0: 294 f_si = f_si[::-1].copy() 295 else: 296 f_i = f_si.sum(0) 297 fm_i = f_si[0] - f_si[1] 298 f_si = np.zeros((4, len(f_i))) 299 f_si[0] = f_i 300 if M > 0: 301 f_si[1:] = M_v[:, np.newaxis] / M * fm_i 302 303 return f_si 304 305 def initialize_from_atomic_densities(self, basis_functions): 306 """Initialize D_asp, nt_sG and Q_aL from atomic densities. 307 308 nt_sG is initialized from atomic orbitals, and will 309 be constructed with the specified magnetic moments and 310 obeying Hund's rules if ``hund`` is true.""" 311 312 # XXX does this work with blacs? What should be distributed? 313 # Apparently this doesn't use blacs at all, so it's serial 314 # with respect to the blacs distribution. That means it works 315 # but is not particularly efficient (not that this is a time 316 # consuming step) 317 318 self.log('Density initialized from atomic densities') 319 320 self.update_atomic_density_matrices( 321 self.setups.empty_atomic_matrix(self.ncomponents, 322 self.atom_partition)) 323 324 f_asi = {} 325 for a in basis_functions.atom_indices: 326 f_asi[a] = self.get_initial_occupations(a) 327 328 # D_asp does not have the same distribution as the basis functions, 329 # so we have to loop over atoms separately. 330 for a in self.D_asp: 331 f_si = f_asi.get(a) 332 if f_si is None: 333 f_si = self.get_initial_occupations(a) 334 self.D_asp[a][:] = self.setups[a].initialize_density_matrix(f_si) 335 336 self.nt_xG = self.gd.zeros(self.ncomponents) 337 basis_functions.add_to_density(self.nt_xG, f_asi) 338 self.nt_sG[:] += self.nct_G 339 self.calculate_normalized_charges_and_mix() 340 341 def initialize_from_wavefunctions(self, wfs): 342 """Initialize D_asp, nt_sG and Q_aL from wave functions.""" 343 self.log('Density initialized from wave functions') 344 self.timer.start('Density initialized from wave functions') 345 self.nt_xG = self.gd.zeros(self.ncomponents) 346 self.calculate_pseudo_density(wfs) 347 self.update_atomic_density_matrices( 348 self.setups.empty_atomic_matrix(self.ncomponents, 349 wfs.atom_partition)) 350 wfs.calculate_atomic_density_matrices(self.D_asp) 351 self.calculate_normalized_charges_and_mix() 352 self.timer.stop('Density initialized from wave functions') 353 354 def initialize_directly_from_arrays(self, nt_sG, nt_vG, D_asp): 355 """Set D_asp and nt_sG directly.""" 356 self.nt_xG = self.gd.zeros(self.ncomponents) 357 self.nt_sG[:] = nt_sG 358 if nt_vG is not None: 359 self.nt_vG[:] = nt_vG 360 361 self.update_atomic_density_matrices(D_asp) 362 D_asp.check_consistency() 363 # No calculate multipole moments? Tests will fail because of 364 # improperly initialized mixer 365 366 def calculate_normalized_charges_and_mix(self): 367 comp_charge, _Q_aL = self.calculate_multipole_moments() 368 self.normalize(comp_charge) 369 self.mix(comp_charge) 370 371 def set_mixer(self, mixer): 372 if mixer is None: 373 mixer = {} 374 if isinstance(mixer, dict): 375 mixer = get_mixer_from_keywords(self.gd.pbc_c.any(), 376 self.ncomponents, **mixer) 377 if not hasattr(mixer, 'mix'): 378 raise ValueError('Not a mixer: %s' % mixer) 379 self.mixer = MixerWrapper(mixer, self.ncomponents, self.gd) 380 381 def estimate_magnetic_moments(self): 382 magmom_av = np.zeros_like(self.magmom_av) 383 magmom_v = np.zeros(3) 384 if self.nspins == 2: 385 for a, D_sp in self.D_asp.items(): 386 M_p = D_sp[0] - D_sp[1] 387 magmom_av[a, 2] = np.dot(M_p, self.setups[a].N0_p) 388 magmom_v[2] += (np.dot(M_p, self.setups[a].Delta_pL[:, 0]) * 389 sqrt(4 * pi)) 390 self.gd.comm.sum(magmom_av) 391 self.gd.comm.sum(magmom_v) 392 magmom_v[2] += self.gd.integrate(self.nt_sG[0] - self.nt_sG[1]) 393 elif not self.collinear: 394 for a, D_sp in self.D_asp.items(): 395 magmom_av[a] = np.dot(D_sp[1:4], self.setups[a].N0_p) 396 magmom_v += (np.dot(D_sp[1:4], self.setups[a].Delta_pL[:, 0]) * 397 sqrt(4 * pi)) 398 # XXXX probably untested code 399 self.gd.comm.sum(magmom_av) 400 self.gd.comm.sum(magmom_v) 401 magmom_v += self.gd.integrate(self.nt_vG) 402 403 return magmom_v, magmom_av 404 405 def get_correction(self, a, spin): 406 """Integrated atomic density correction. 407 408 Get the integrated correction to the pseuso density relative to 409 the all-electron density. 410 """ 411 setup = self.setups[a] 412 return sqrt(4 * pi) * ( 413 np.dot(self.D_asp[a][spin], setup.Delta_pL[:, 0]) + 414 setup.Delta0 / self.nspins) 415 416 def get_all_electron_density(self, atoms=None, gridrefinement=2, 417 spos_ac=None, skip_core=False): 418 """Return real all-electron density array. 419 420 Usage: Either get_all_electron_density(atoms) or 421 get_all_electron_density(spos_ac=spos_ac) 422 423 skip_core=True theoretically returns the 424 all-electron valence density (use with 425 care; will not in general integrate 426 to valence) 427 """ 428 if spos_ac is None: 429 spos_ac = atoms.get_scaled_positions() % 1.0 430 431 # Refinement of coarse grid, for representation of the AE-density 432 # XXXXXXXXXXXX think about distribution depending on gridrefinement! 433 if gridrefinement == 1: 434 gd = self.redistributor.aux_gd 435 n_sg = self.nt_sG.copy() 436 # This will get the density with the same distribution 437 # as finegd: 438 n_sg = self.redistributor.distribute(n_sg) 439 elif gridrefinement == 2: 440 gd = self.finegd 441 if self.nt_sg is None: 442 self.interpolate_pseudo_density() 443 n_sg = self.nt_sg.copy() 444 elif gridrefinement == 4: 445 # Extra fine grid 446 gd = self.finegd.refine() 447 448 # Interpolation function for the density: 449 interpolator = Transformer(self.finegd, gd, 3) # XXX grids! 450 451 # Transfer the pseudo-density to the fine grid: 452 n_sg = gd.empty(self.nspins) 453 if self.nt_sg is None: 454 self.interpolate_pseudo_density() 455 for s in range(self.nspins): 456 interpolator.apply(self.nt_sg[s], n_sg[s]) 457 else: 458 raise NotImplementedError 459 460 # Add corrections to pseudo-density to get the AE-density 461 splines = {} 462 phi_aj = [] 463 phit_aj = [] 464 nc_a = [] 465 nct_a = [] 466 for a, id in enumerate(self.setups.id_a): 467 if id in splines: 468 phi_j, phit_j, nc, nct = splines[id] 469 else: 470 # Load splines: 471 phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4] 472 splines[id] = (phi_j, phit_j, nc, nct) 473 phi_aj.append(phi_j) 474 phit_aj.append(phit_j) 475 nc_a.append([nc]) 476 nct_a.append([nct]) 477 478 # Create localized functions from splines 479 phi = BasisFunctions(gd, phi_aj) 480 phit = BasisFunctions(gd, phit_aj) 481 nc = LFC(gd, nc_a) 482 nct = LFC(gd, nct_a) 483 phi.set_positions(spos_ac) 484 phit.set_positions(spos_ac) 485 nc.set_positions(spos_ac) 486 nct.set_positions(spos_ac) 487 488 I_sa = np.zeros((self.nspins, len(spos_ac))) 489 a_W = np.empty(len(phi.M_W), np.intc) 490 W = 0 491 for a in phi.atom_indices: 492 nw = len(phi.sphere_a[a].M_w) 493 a_W[W:W + nw] = a 494 W += nw 495 496 x_W = phi.create_displacement_arrays()[0] 497 498 # We need the charges for the density matrices in order to add 499 # nuclear charges at each atom. Hence we use the aux partition: 500 # The one where atoms are distributed according to which realspace 501 # domain they belong to. 502 D_asp = self.atomdist.to_aux(self.D_asp) 503 504 rho_MM = np.zeros((phi.Mmax, phi.Mmax)) 505 for s, I_a in enumerate(I_sa): 506 M1 = 0 507 for a, setup in enumerate(self.setups): 508 ni = setup.ni 509 D_sp = D_asp.get(a) 510 if D_sp is None: 511 D_sp = np.empty((self.nspins, ni * (ni + 1) // 2)) 512 else: 513 I_a[a] = ((setup.Nct) / self.nspins - 514 sqrt(4 * pi) * 515 np.dot(D_sp[s], setup.Delta_pL[:, 0])) 516 517 if not skip_core: 518 I_a[a] -= setup.Nc / self.nspins 519 520 rank = D_asp.partition.rank_a[a] 521 D_asp.partition.comm.broadcast(D_sp, rank) 522 M2 = M1 + ni 523 rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s]) 524 M1 = M2 525 526 assert np.all(n_sg[s].shape == phi.gd.n_c) 527 phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a, 528 x_W) 529 phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a, 530 x_W) 531 532 # wth is this? 533 a_W = np.empty(len(nc.M_W), np.intc) 534 W = 0 535 for a in nc.atom_indices: 536 nw = len(nc.sphere_a[a].M_w) 537 a_W[W:W + nw] = a 538 W += nw 539 scale = 1.0 / self.nspins 540 541 for s, I_a in enumerate(I_sa): 542 543 if not skip_core: 544 nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a) 545 546 nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a) 547 D_asp.partition.comm.sum(I_a) 548 N_c = gd.N_c 549 g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c 550 551 if not skip_core: 552 for I, g_c in zip(I_a, g_ac): 553 if (g_c >= 0).all() and (g_c < gd.n_c).all(): 554 n_sg[s][tuple(g_c)] -= I / gd.dv 555 556 return n_sg, gd 557 558 def estimate_memory(self, mem): 559 nspins = self.nspins 560 nbytes = self.gd.bytecount() 561 nfinebytes = self.finegd.bytecount() 562 563 arrays = mem.subnode('Arrays') 564 for name, size in [('nt_sG', nbytes * nspins), 565 ('nt_sg', nfinebytes * nspins), 566 ('nt_g', nfinebytes), 567 ('rhot_g', nfinebytes), 568 ('nct_G', nbytes)]: 569 arrays.subnode(name, size) 570 571 lfs = mem.subnode('Localized functions') 572 for name, obj in [('nct', self.nct), 573 ('ghat', self.ghat)]: 574 obj.estimate_memory(lfs.subnode(name)) 575 self.mixer.estimate_memory(mem.subnode('Mixer'), self.gd) 576 577 # TODO 578 # The implementation of interpolator memory use is not very 579 # accurate; 20 MiB vs 13 MiB estimated in one example, probably 580 # worse for parallel calculations. 581 582 def get_spin_contamination(self, atoms, majority_spin=0): 583 """Calculate the spin contamination. 584 585 Spin contamination is defined as the integral over the 586 spin density difference, where it is negative (i.e. the 587 minority spin density is larger than the majority spin density. 588 """ 589 590 if majority_spin == 0: 591 smaj = 0 592 smin = 1 593 else: 594 smaj = 1 595 smin = 0 596 nt_sg, gd = self.get_all_electron_density(atoms) 597 dt_sg = nt_sg[smin] - nt_sg[smaj] 598 dt_sg = np.where(dt_sg > 0, dt_sg, 0.0) 599 return gd.integrate(dt_sg) 600 601 def write(self, writer): 602 writer.write(density=self.gd.collect(self.nt_xG) / Bohr**3, 603 atomic_density_matrices=pack_atomic_matrices(self.D_asp)) 604 605 def read(self, reader): 606 nt_xG = self.gd.empty(self.ncomponents) 607 self.gd.distribute(reader.density.density, nt_xG) 608 nt_xG *= reader.bohr**3 609 610 # Read atomic density matrices 611 natoms = len(self.setups) 612 atom_partition = AtomPartition(self.gd.comm, np.zeros(natoms, int), 613 'density-gd') 614 D_asp = self.setups.empty_atomic_matrix(self.ncomponents, 615 atom_partition) 616 self.atom_partition = atom_partition # XXXXXX 617 spos_ac = np.zeros((natoms, 3)) # XXXX 618 self.atomdist = self.redistributor.get_atom_distributions(spos_ac) 619 620 D_sP = reader.density.atomic_density_matrices 621 if self.gd.comm.rank == 0: 622 D_asp.update(unpack_atomic_matrices(D_sP, self.setups)) 623 D_asp.check_consistency() 624 625 if self.collinear: 626 nt_sG = nt_xG 627 nt_vG = None 628 else: 629 nt_sG = nt_xG[:1] 630 nt_vG = nt_xG[1:] 631 632 self.initialize_directly_from_arrays(nt_sG, nt_vG, D_asp) 633 634 def initialize_from_other_density(self, dens, kptband_comm): 635 """Redistribute pseudo density and atomic density matrices. 636 637 Collect dens.nt_sG and dens.D_asp to world master and distribute.""" 638 639 new_nt_sG = redistribute_array(dens.nt_sG, dens.gd, self.gd, 640 self.nspins, kptband_comm) 641 self.atom_partition, self.atomdist = \ 642 create_atom_partition_and_distibutions(self.gd, self.nspins, 643 self.setups, 644 self.redistributor, 645 kptband_comm) 646 D_asp = \ 647 redistribute_atomic_matrices(dens.D_asp, self.gd, self.nspins, 648 self.setups, self.atom_partition, 649 kptband_comm) 650 self.initialize_directly_from_arrays(new_nt_sG, None, D_asp) 651 652 653class RealSpaceDensity(Density): 654 def __init__(self, gd, finegd, nspins, collinear, charge, redistributor, 655 stencil=3, 656 background_charge=None): 657 Density.__init__(self, gd, finegd, nspins, collinear, 658 charge, redistributor, 659 background_charge=background_charge) 660 self.stencil = stencil 661 662 self.interpolator = None 663 664 def initialize(self, setups, timer, magmom_a, hund): 665 Density.initialize(self, setups, timer, magmom_a, hund) 666 667 # Interpolation function for the density: 668 self.interpolator = Transformer(self.redistributor.aux_gd, 669 self.finegd, self.stencil) 670 671 spline_aj = [] 672 for setup in setups: 673 if setup.nct is None: 674 spline_aj.append([]) 675 else: 676 spline_aj.append([setup.nct]) 677 self.nct = LFC(self.gd, spline_aj, 678 integral=[setup.Nct for setup in setups], 679 forces=True, cut=True) 680 self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups], 681 integral=sqrt(4 * pi), forces=True) 682 683 def set_positions(self, spos_ac, atom_partition): 684 Density.set_positions(self, spos_ac, atom_partition) 685 self.nct_G = self.gd.zeros() 686 self.nct.add(self.nct_G, 1.0 / self.nspins) 687 688 def interpolate_pseudo_density(self, comp_charge=None): 689 """Interpolate pseudo density to fine grid.""" 690 if comp_charge is None: 691 comp_charge, _Q_aL = self.calculate_multipole_moments() 692 693 self.nt_sg = self.distribute_and_interpolate(self.nt_sG, self.nt_sg) 694 695 # With periodic boundary conditions, the interpolation will 696 # conserve the number of electrons. 697 if not self.gd.pbc_c.all(): 698 # With zero-boundary conditions in one or more directions, 699 # this is not the case. 700 pseudo_charge = (self.background_charge.charge - self.charge - 701 comp_charge) 702 if abs(pseudo_charge) > 1.0e-14: 703 x = (pseudo_charge / 704 self.finegd.integrate(self.nt_sg).sum()) 705 self.nt_sg *= x 706 707 def interpolate(self, in_xR, out_xR=None): 708 """Interpolate array(s).""" 709 710 # ndim will be 3 in finite-difference mode and 1 when working 711 # with the AtomPAW class (spherical atoms and 1d grids) 712 ndim = self.gd.ndim 713 714 if out_xR is None: 715 out_xR = self.finegd.empty(in_xR.shape[:-ndim]) 716 717 a_xR = in_xR.reshape((-1,) + in_xR.shape[-ndim:]) 718 b_xR = out_xR.reshape((-1,) + out_xR.shape[-ndim:]) 719 720 for in_R, out_R in zip(a_xR, b_xR): 721 self.interpolator.apply(in_R, out_R) 722 723 return out_xR 724 725 def distribute_and_interpolate(self, in_xR, out_xR=None): 726 in_xR = self.redistributor.distribute(in_xR) 727 return self.interpolate(in_xR, out_xR) 728 729 def calculate_pseudo_charge(self): 730 self.nt_g = self.nt_sg.sum(axis=0) 731 self.rhot_g = self.nt_g.copy() 732 self.calculate_multipole_moments() 733 self.ghat.add(self.rhot_g, self.Q_aL) 734 self.background_charge.add_charge_to(self.rhot_g) 735 736 if debug: 737 charge = self.finegd.integrate(self.rhot_g) + self.charge 738 if abs(charge) > self.charge_eps: 739 raise RuntimeError('Charge not conserved: excess=%.9f' % 740 charge) 741 742 def get_pseudo_core_kinetic_energy_density_lfc(self): 743 return LFC(self.gd, 744 [[setup.tauct] for setup in self.setups], 745 forces=True, cut=True) 746 747 def calculate_dipole_moment(self): 748 return self.finegd.calculate_dipole_moment(self.rhot_g) 749 750 751def redistribute_array(nt_sG, gd1, gd2, nspins, kptband_comm): 752 nt_sG = gd1.collect(nt_sG) 753 new_nt_sG = gd2.empty(nspins) 754 if kptband_comm.rank == 0: 755 gd2.distribute(nt_sG, new_nt_sG) 756 kptband_comm.broadcast(new_nt_sG, 0) 757 return new_nt_sG 758 759 760def create_atom_partition_and_distibutions(gd2, nspins, setups, redistributor, 761 kptband_comm): 762 natoms = len(setups) 763 atom_partition = AtomPartition(gd2.comm, np.zeros(natoms, int), 764 'density-gd') 765 spos_ac = np.zeros((natoms, 3)) # XXXX 766 atomdist = redistributor.get_atom_distributions(spos_ac) 767 return atom_partition, atomdist 768 769 770def redistribute_atomic_matrices(D_asp, gd2, nspins, setups, atom_partition, 771 kptband_comm): 772 D_sP = pack_atomic_matrices(D_asp) 773 D_asp = setups.empty_atomic_matrix(nspins, atom_partition) 774 if gd2.comm.rank == 0: 775 if kptband_comm.rank > 0: 776 nP = sum(setup.ni * (setup.ni + 1) // 2 777 for setup in setups) 778 D_sP = np.empty((nspins, nP)) 779 kptband_comm.broadcast(D_sP, 0) 780 D_asp.update(unpack_atomic_matrices(D_sP, setups)) 781 D_asp.check_consistency() 782 return D_asp 783