Coverage for gpaw/density.py: 97%

1# Copyright (C) 2003 CAMP 

2# Please see the accompanying LICENSE file for further information. 

3 

4"""This module defines a density class.""" 

5 

6from math import pi, sqrt 

7 

8import numpy as np 

9from ase.units import Bohr 

10 

11from gpaw import debug 

12from gpaw.mixer import get_mixer_from_keywords, MixerWrapper 

13from gpaw.transformers import Transformer 

14from gpaw.lfc import LFC, BasisFunctions 

15from gpaw.wavefunctions.lcao import LCAOWaveFunctions 

16from gpaw.utilities import (unpack_density, unpack_atomic_matrices, 

17 pack_atomic_matrices) 

18from gpaw.utilities.partition import AtomPartition 

19from gpaw.utilities.timing import nulltimer 

20from gpaw.arraydict import ArrayDict 

21 

22 

23class CompensationChargeExpansionCoefficients: 

24 def __init__(self, setups, nspins): 

25 self.setups = setups 

26 self.nspins = nspins 

27 

28 def calculate(self, D_asp): 

29 """Calculate multipole moments of compensation charges. 

30 

31 Returns the total compensation charge in units of electron 

32 charge, so the number will be negative because of the 

33 dominating contribution from the nuclear charge.""" 

34 atom_partition = D_asp.partition 

35 shape_a = [(setup.Delta_pL.shape[1],) for setup in self.setups] 

36 Q_aL = atom_partition.arraydict(shape_a, dtype=float) 

37 for a, D_sp in D_asp.items(): 

38 setup = self.setups[a] 

39 Q_L = np.dot(D_sp[:self.nspins].sum(0), setup.Delta_pL) 

40 Q_L[0] += setup.Delta0 

41 Q_aL[a] = Q_L 

42 return Q_aL 

43 

44 def get_charge(self, Q_aL): 

45 local_charge = sqrt(4 * pi) * sum(Q_L[0] for Q_L in Q_aL.values()) 

46 return Q_aL.partition.comm.sum_scalar(local_charge) 

47 

48 

49class NullBackgroundCharge: 

50 charge = 0.0 

51 

52 def set_grid_descriptor(self, gd): 

53 pass 

54 

55 def add_charge_to(self, rhot_g): 

56 pass 

57 

58 def add_fourier_space_charge_to(self, pd, rhot_q): 

59 pass 

60 

61 

62class Density: 

63 """Density object. 

64 

65 Attributes: 

66 =============== ===================================================== 

67 ``gd`` Grid descriptor for coarse grids. 

68 ``finegd`` Grid descriptor for fine grids. 

69 ``interpolate`` Function for interpolating the electron density. 

70 ``mixer`` ``DensityMixer`` object. 

71 =============== ===================================================== 

72 

73 Soft and smooth pseudo functions on uniform 3D grids: 

74 ========== ========================================= 

75 ``nt_sG`` Electron density on the coarse grid. 

76 ``nt_sg`` Electron density on the fine grid. 

77 ``nt_g`` Electron density on the fine grid. 

78 ``rhot_g`` Charge density on the fine grid. 

79 ``nct_G`` Core electron-density on the coarse grid. 

80 ========== ========================================= 

81 """ 

82 

83 def __init__(self, gd, finegd, nspins, collinear, charge, redistributor, 

84 background_charge=None): 

85 """Create the Density object.""" 

86 

87 self.gd = gd 

88 self.finegd = finegd 

89 self.nspins = nspins 

90 self.collinear = collinear 

91 self.charge = float(charge) 

92 self.redistributor = redistributor 

93 self.atomdist = None 

94 

95 self.ncomponents = self.nspins if self.collinear else 1 + 3 

96 

97 # This can contain e.g. a Jellium background charge 

98 if background_charge is None: 

99 background_charge = NullBackgroundCharge() 

100 background_charge.set_grid_descriptor(self.finegd) 

101 self.background_charge = background_charge 

102 

103 self.charge_eps = 1e-7 

104 

105 self.D_asp = None 

106 self.Q = CompensationChargeExpansionCoefficients([], self.nspins) 

107 self.Q_aL = None 

108 

109 self.nct_G = None 

110 self.nt_xG = None 

111 self.rhot_g = None 

112 self.nt_xg = None 

113 self.nt_sg = None 

114 self.nt_vg = None 

115 self.nt_g = None 

116 

117 self.atom_partition = None 

118 

119 self.setups = None 

120 self.hund = None 

121 self.magmom_av = None 

122 

123 self.fixed = False 

124 # XXX at least one test will fail because None has no 'reset()' 

125 # So we need DummyMixer I guess 

126 self.mixer = None 

127 self.set_mixer(None) 

128 

129 self.timer = nulltimer 

130 self.error = None 

131 self.nct = None 

132 self.ghat = None 

133 self.log = None 

134 

135 @property 

136 def nt_sG(self): 

137 return None if self.nt_xG is None else self.nt_xG[:self.nspins] 

138 

139 @property 

140 def nt_vG(self): 

141 return None if self.nt_xG is None else self.nt_xG[self.nspins:] 

142 

143 def __str__(self): 

144 s = 'Densities:\n' 

145 s += ' Coarse grid: {}*{}*{} grid\n'.format(*self.gd.N_c) 

146 s += ' Fine grid: {}*{}*{} grid\n'.format(*self.finegd.N_c) 

147 s += f' Total Charge: {self.charge:.6f}' 

148 if self.fixed: 

149 s += '\n Fixed' 

150 return s 

151 

152 def summary(self, atoms, magmom, log): 

153 if self.nspins == 1: 

154 return 

155 try: 

156 # XXX This doesn't always work, HGH, SIC, ... 

157 sc = self.get_spin_contamination(atoms, int(magmom < 0)) 

158 log('Spin contamination: %f electrons' % sc) 

159 except (TypeError, AttributeError, AssertionError): 

160 pass 

161 

162 def initialize(self, setups, timer, magmom_av, hund): 

163 self.timer = timer 

164 self.setups = setups 

165 self.Q = CompensationChargeExpansionCoefficients(setups, self.nspins) 

166 self.hund = hund 

167 self.magmom_av = magmom_av 

168 

169 def reset(self): 

170 # TODO: reset other parameters? 

171 self.nt_xG = None 

172 

173 def set_positions_without_ruining_everything(self, spos_ac, 

174 atom_partition): 

175 rank_a = atom_partition.rank_a 

176 # If both old and new atomic ranks are present, start a blank dict if 

177 # it previously didn't exist but it will needed for the new atoms. 

178 if (self.atom_partition is not None and 

179 self.D_asp is None and (rank_a == self.gd.comm.rank).any()): 

180 self.update_atomic_density_matrices( 

181 self.setups.empty_atomic_matrix(self.ncomponents, 

182 self.atom_partition)) 

183 

184 if (self.atom_partition is not None and self.D_asp is not None and 

185 self.gd.comm.size > 1): 

186 self.timer.start('Redistribute') 

187 self.D_asp.redistribute(atom_partition) 

188 self.timer.stop('Redistribute') 

189 

190 self.atom_partition = atom_partition 

191 self.atomdist = self.redistributor.get_atom_distributions(spos_ac) 

192 

193 def set_positions(self, spos_ac, atom_partition): 

194 self.set_positions_without_ruining_everything(spos_ac, atom_partition) 

195 self.nct.set_positions(spos_ac, atom_partition) 

196 self.ghat.set_positions(spos_ac, atom_partition) 

197 self.mixer.reset() 

198 

199 self.nt_xg = None 

200 self.nt_sg = None 

201 self.nt_vg = None 

202 self.nt_g = None 

203 self.rhot_g = None 

204 

205 def calculate_pseudo_density(self, wfs): 

206 """Calculate nt_sG from scratch. 

207 

208 nt_sG will be equal to nct_G plus the contribution from 

209 wfs.add_to_density(). 

210 """ 

211 wfs.calculate_density_contribution(self.nt_xG) 

212 self.nt_sG[:] += self.nct_G 

213 

214 def update_atomic_density_matrices(self, value): 

215 if isinstance(value, dict): 

216 tmp = self.setups.empty_atomic_matrix(self.ncomponents, 

217 self.atom_partition) 

218 tmp.update(value) 

219 value = tmp 

220 assert isinstance(value, ArrayDict) or value is None, type(value) 

221 if value is not None: 

222 value.check_consistency() 

223 self.D_asp = value 

224 

225 def update(self, wfs): 

226 self.timer.start('Density') 

227 with self.timer('Pseudo density'): 

228 self.calculate_pseudo_density(wfs) 

229 with self.timer('Atomic density matrices'): 

230 wfs.calculate_atomic_density_matrices(self.D_asp) 

231 with self.timer('Multipole moments'): 

232 comp_charge, _Q_aL = self.calculate_multipole_moments() 

233 

234 if isinstance(wfs, LCAOWaveFunctions): 

235 self.timer.start('Normalize') 

236 self.normalize(comp_charge) 

237 self.timer.stop('Normalize') 

238 

239 self.timer.start('Mix') 

240 self.mix(comp_charge) 

241 self.timer.stop('Mix') 

242 self.timer.stop('Density') 

243 

244 def normalize(self, comp_charge): 

245 """Normalize pseudo density.""" 

246 

247 pseudo_charge = self.gd.integrate(self.nt_sG).sum() 

248 

249 if (pseudo_charge + self.charge + comp_charge - 

250 self.background_charge.charge != 0): 

251 if pseudo_charge != 0: 

252 x = (self.background_charge.charge - self.charge - 

253 comp_charge) / pseudo_charge 

254 self.nt_xG *= x 

255 else: 

256 # Use homogeneous background. 

257 # 

258 # We have to use the volume per actual grid point, 

259 # which is not the same as gd.volume as the latter 

260 # includes ghost points. 

261 volume = self.gd.get_size_of_global_array().prod() * self.gd.dv 

262 total_charge = (self.charge + comp_charge - 

263 self.background_charge.charge) 

264 self.nt_sG[:] = -total_charge / volume 

265 

266 def mix(self, comp_charge): 

267 assert isinstance(self.mixer, MixerWrapper), self.mixer 

268 self.error = self.mixer.mix(self.nt_xG, self.D_asp) 

269 assert self.error is not None, self.mixer 

270 

271 comp_charge = None 

272 self.interpolate_pseudo_density(comp_charge) 

273 self.calculate_pseudo_charge() 

274 

275 def calculate_multipole_moments(self): 

276 D_asp = self.atomdist.to_aux(self.D_asp) 

277 Q_aL = self.Q.calculate(D_asp) 

278 self.Q_aL = Q_aL 

279 return self.Q.get_charge(Q_aL), Q_aL 

280 

281 def get_initial_occupations(self, a): 

282 # distribute charge on all atoms 

283 # XXX interaction with background charge may be finicky 

284 c = (self.charge - self.background_charge.charge) / len(self.setups) 

285 M_v = self.magmom_av[a] 

286 M = np.linalg.norm(M_v) 

287 f_si = self.setups[a].calculate_initial_occupation_numbers( 

288 M, self.hund, charge=c, 

289 nspins=self.nspins if self.collinear else 2) 

290 

291 if self.collinear: 

292 if M_v[2] < 0: 

293 f_si = f_si[::-1].copy() 

294 else: 

295 f_i = f_si.sum(0) 

296 fm_i = f_si[0] - f_si[1] 

297 f_si = np.zeros((4, len(f_i))) 

298 f_si[0] = f_i 

299 if M > 0: 

300 f_si[1:] = M_v[:, np.newaxis] / M * fm_i 

301 

302 return f_si 

303 

304 def initialize_from_atomic_densities(self, basis_functions): 

305 """Initialize D_asp, nt_sG and Q_aL from atomic densities. 

306 

307 nt_sG is initialized from atomic orbitals, and will 

308 be constructed with the specified magnetic moments and 

309 obeying Hund's rules if ``hund`` is true.""" 

310 

311 # XXX does this work with blacs? What should be distributed? 

312 # Apparently this doesn't use blacs at all, so it's serial 

313 # with respect to the blacs distribution. That means it works 

314 # but is not particularly efficient (not that this is a time 

315 # consuming step) 

316 

317 self.log('Density initialized from atomic densities') 

318 

319 self.update_atomic_density_matrices( 

320 self.setups.empty_atomic_matrix(self.ncomponents, 

321 self.atom_partition)) 

322 

323 f_asi = {} 

324 for a in basis_functions.atom_indices: 

325 f_asi[a] = self.get_initial_occupations(a) 

326 

327 # D_asp does not have the same distribution as the basis functions, 

328 # so we have to loop over atoms separately. 

329 for a in self.D_asp: 

330 f_si = f_asi.get(a) 

331 if f_si is None: 

332 f_si = self.get_initial_occupations(a) 

333 self.D_asp[a][:] = self.setups[a].initialize_density_matrix(f_si) 

334 

335 self.nt_xG = self.gd.zeros(self.ncomponents) 

336 basis_functions.add_to_density(self.nt_xG, f_asi) 

337 self.nt_sG[:] += self.nct_G 

338 self.calculate_normalized_charges_and_mix() 

339 

340 def initialize_from_wavefunctions(self, wfs): 

341 """Initialize D_asp, nt_sG and Q_aL from wave functions.""" 

342 self.log('Density initialized from wave functions') 

343 self.timer.start('Density initialized from wave functions') 

344 self.nt_xG = self.gd.zeros(self.ncomponents) 

345 self.calculate_pseudo_density(wfs) 

346 self.update_atomic_density_matrices( 

347 self.setups.empty_atomic_matrix(self.ncomponents, 

348 wfs.atom_partition)) 

349 wfs.calculate_atomic_density_matrices(self.D_asp) 

350 self.calculate_normalized_charges_and_mix() 

351 self.timer.stop('Density initialized from wave functions') 

352 

353 def initialize_directly_from_arrays(self, nt_sG, nt_vG, D_asp): 

354 """Set D_asp and nt_sG directly.""" 

355 self.nt_xG = self.gd.zeros(self.ncomponents) 

356 self.nt_sG[:] = nt_sG 

357 if nt_vG is not None: 

358 self.nt_vG[:] = nt_vG 

359 

360 self.update_atomic_density_matrices(D_asp) 

361 D_asp.check_consistency() 

362 # No calculate multipole moments? Tests will fail because of 

363 # improperly initialized mixer 

364 

365 def calculate_normalized_charges_and_mix(self): 

366 comp_charge, _Q_aL = self.calculate_multipole_moments() 

367 self.normalize(comp_charge) 

368 self.mix(comp_charge) 

369 

370 def set_mixer(self, mixer): 

371 if mixer is None: 

372 mixer = {} 

373 if isinstance(mixer, dict): 

374 mixer = get_mixer_from_keywords(self.gd.pbc_c.any(), 

375 self.ncomponents, **mixer) 

376 if not hasattr(mixer, 'mix'): 

377 raise ValueError('Not a mixer: %s' % mixer) 

378 self.mixer = MixerWrapper(mixer, self.ncomponents, self.gd) 

379 

380 def calculate_magnetic_moments(self): 

381 magmom_av = np.zeros_like(self.magmom_av) 

382 magmom_v = np.zeros(3) 

383 if self.nspins == 2: 

384 for a, D_sp in self.D_asp.items(): 

385 M_p = D_sp[0] - D_sp[1] 

386 magmom_av[a, 2] = np.dot(M_p, self.setups[a].N0_p) 

387 magmom_v[2] += (np.dot(M_p, self.setups[a].Delta_pL[:, 0]) * 

388 sqrt(4 * pi)) 

389 self.gd.comm.sum(magmom_av) 

390 self.gd.comm.sum(magmom_v) 

391 magmom_v[2] += self.gd.integrate(self.nt_sG[0] - self.nt_sG[1]) 

392 elif not self.collinear: 

393 for a, D_sp in self.D_asp.items(): 

394 magmom_av[a] = np.dot(D_sp[1:4], self.setups[a].N0_p) 

395 magmom_v += (np.dot(D_sp[1:4], self.setups[a].Delta_pL[:, 0]) * 

396 sqrt(4 * pi)) 

397 # XXXX probably untested code 

398 self.gd.comm.sum(magmom_av) 

399 self.gd.comm.sum(magmom_v) 

400 magmom_v += self.gd.integrate(self.nt_vG) 

401 

402 return magmom_v, magmom_av 

403 

404 estimate_magnetic_moments = calculate_magnetic_moments 

405 

406 def get_correction(self, a, spin): 

407 """Integrated atomic density correction. 

408 

409 Get the integrated correction to the pseuso density relative to 

410 the all-electron density. 

411 """ 

412 setup = self.setups[a] 

413 return sqrt(4 * pi) * ( 

414 np.dot(self.D_asp[a][spin], setup.Delta_pL[:, 0]) + 

415 setup.Delta0 / self.nspins) 

416 

417 def get_all_electron_density(self, atoms=None, gridrefinement=2, 

418 spos_ac=None, skip_core=False): 

419 """Return real all-electron density array. 

420 

421 Usage: Either get_all_electron_density(atoms) or 

422 get_all_electron_density(spos_ac=spos_ac) 

423 

424 skip_core=True theoretically returns the 

425 all-electron valence density (use with 

426 care; will not in general integrate 

427 to valence) 

428 """ 

429 if spos_ac is None: 

430 spos_ac = atoms.get_scaled_positions() % 1.0 

431 spos_ch = [] # for coreholes 

432 

433 # Refinement of coarse grid, for representation of the AE-density 

434 # XXXXXXXXXXXX think about distribution depending on gridrefinement! 

435 if gridrefinement == 1: 

436 gd = self.redistributor.aux_gd 

437 n_sg = self.nt_sG.copy() 

438 # This will get the density with the same distribution 

439 # as finegd: 

440 n_sg = self.redistributor.distribute(n_sg) 

441 elif gridrefinement == 2: 

442 gd = self.finegd 

443 if self.nt_sg is None: 

444 self.interpolate_pseudo_density() 

445 n_sg = self.nt_sg.copy() 

446 elif gridrefinement == 4: 

447 # Extra fine grid 

448 gd = self.finegd.refine() 

449 

450 # Interpolation function for the density: 

451 interpolator = Transformer(self.finegd, gd, 3) # XXX grids! 

452 

453 # Transfer the pseudo-density to the fine grid: 

454 n_sg = gd.empty(self.nspins) 

455 if self.nt_sg is None: 

456 self.interpolate_pseudo_density() 

457 for s in range(self.nspins): 

458 interpolator.apply(self.nt_sg[s], n_sg[s]) 

459 else: 

460 raise NotImplementedError 

461 

462 if self.nspins > 1 and not skip_core: 

463 # Support for corehole in spin-polarized system 

464 spos_ch = [] 

465 n_ch_a = [] 

466 n_ch = [] 

467 

468 # Add corrections to pseudo-density to get the AE-density 

469 splines = {} 

470 phi_aj = [] 

471 phit_aj = [] 

472 nc_a = [] 

473 nct_a = [] 

474 for a, id in enumerate(self.setups.id_a): 

475 if id in splines: 

476 phi_j, phit_j, nc, nct = splines[id] 

477 else: 

478 # Load splines: 

479 phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4] 

480 splines[id] = (phi_j, phit_j, nc, nct) 

481 phi_aj.append(phi_j) 

482 phit_aj.append(phit_j) 

483 nc_a.append([nc]) 

484 nct_a.append([nct]) 

485 if self.setups[a].data.has_corehole and not skip_core and \ 

486 self.nspins > 1: 

487 work_setup = self.setups[a].data 

488 rmax = nc.get_cutoff() 

489 # work_setup.phicorehole_g 

490 phi_ch = work_setup.phicorehole_g 

491 phi_ch = np.where(abs(phi_ch) < 1e-160, 0, phi_ch) 

492 n_ch = np.dot(work_setup.fcorehole, phi_ch**2) / (4 * pi) 

493 # n_ch[0] = n_ch[1] # ch is allready taylored 

494 # fcorehole should divided by two - use scale for this 

495 n_ch_spl = work_setup.rgd.spline(n_ch, rmax, points=1000) 

496 n_ch_a.append([n_ch_spl]) 

497 spos_ch.append(spos_ac[a]) 

498 

499 # Create localized functions from splines 

500 phi = BasisFunctions(gd, phi_aj) 

501 phit = BasisFunctions(gd, phit_aj) 

502 nc = LFC(gd, nc_a) 

503 nct = LFC(gd, nct_a) 

504 phi.set_positions(spos_ac) 

505 phit.set_positions(spos_ac) 

506 nc.set_positions(spos_ac) 

507 nct.set_positions(spos_ac) 

508 if spos_ch: 

509 nch = LFC(gd, n_ch_a) 

510 nch.set_positions(spos_ch) 

511 

512 I_sa = np.zeros((self.nspins, len(spos_ac))) 

513 a_W = np.empty(len(phi.M_W), np.intc) 

514 W = 0 

515 for a in phi.atom_indices: 

516 nw = len(phi.sphere_a[a].M_w) 

517 a_W[W:W + nw] = a 

518 W += nw 

519 

520 x_W = phi.create_displacement_arrays()[0] 

521 

522 # We need the charges for the density matrices in order to add 

523 # nuclear charges at each atom. Hence we use the aux partition: 

524 # The one where atoms are distributed according to which realspace 

525 # domain they belong to. 

526 D_asp = self.atomdist.to_aux(self.D_asp) 

527 

528 rho_MM = np.zeros((phi.Mmax, phi.Mmax)) 

529 for s, I_a in enumerate(I_sa): 

530 M1 = 0 

531 for a, setup in enumerate(self.setups): 

532 ni = setup.ni 

533 D_sp = D_asp.get(a) 

534 if D_sp is None: 

535 D_sp = np.empty((self.nspins, ni * (ni + 1) // 2)) 

536 else: 

537 I_a[a] = ((setup.Nct) / self.nspins - 

538 sqrt(4 * pi) * 

539 np.dot(D_sp[s], setup.Delta_pL[:, 0])) 

540 

541 if not skip_core: 

542 I_a[a] -= setup.Nc / self.nspins 

543 if self.setups[a].data.has_corehole and \ 

544 self.nspins > 1: 

545 I_a[a] += pow(-1, s) * \ 

546 self.setups[a].data.fcorehole / 2 

547 

548 rank = D_asp.partition.rank_a[a] 

549 D_asp.partition.comm.broadcast(D_sp, rank) 

550 M2 = M1 + ni 

551 rho_MM[M1:M2, M1:M2] = unpack_density(D_sp[s]) 

552 M1 = M2 

553 

554 assert np.all(n_sg[s].shape == phi.gd.n_c) 

555 phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a, 

556 x_W) 

557 phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a, 

558 x_W) 

559 

560 # wth is this? 

561 a_W = np.empty(len(nc.M_W), np.intc) 

562 W = 0 

563 for a in nc.atom_indices: 

564 nw = len(nc.sphere_a[a].M_w) 

565 a_W[W:W + nw] = a 

566 W += nw 

567 scale = 1.0 / self.nspins 

568 

569 for s, I_a in enumerate(I_sa): 

570 

571 if not skip_core: 

572 nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a) 

573 # correct for ch here 

574 if spos_ch: 

575 nch.lfc.ae_core_density_correction(- scale * pow(-1, s), 

576 n_sg[s], a_W, I_a) 

577 nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a) 

578 D_asp.partition.comm.sum(I_a) 

579 N_c = gd.N_c 

580 g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c 

581 

582 if not skip_core: 

583 for I, g_c in zip(I_a, g_ac): 

584 if (g_c >= 0).all() and (g_c < gd.n_c).all(): 

585 n_sg[s][tuple(g_c)] -= I / gd.dv 

586 

587 return n_sg, gd 

588 

589 def estimate_memory(self, mem): 

590 nspins = self.nspins 

591 nbytes = self.gd.bytecount() 

592 nfinebytes = self.finegd.bytecount() 

593 

594 arrays = mem.subnode('Arrays') 

595 for name, size in [('nt_sG', nbytes * nspins), 

596 ('nt_sg', nfinebytes * nspins), 

597 ('nt_g', nfinebytes), 

598 ('rhot_g', nfinebytes), 

599 ('nct_G', nbytes)]: 

600 arrays.subnode(name, size) 

601 

602 lfs = mem.subnode('Localized functions') 

603 for name, obj in [('nct', self.nct), 

604 ('ghat', self.ghat)]: 

605 obj.estimate_memory(lfs.subnode(name)) 

606 self.mixer.estimate_memory(mem.subnode('Mixer'), self.gd) 

607 

608 # TODO 

609 # The implementation of interpolator memory use is not very 

610 # accurate; 20 MiB vs 13 MiB estimated in one example, probably 

611 # worse for parallel calculations. 

612 

613 def get_spin_contamination(self, atoms, majority_spin=0): 

614 """Calculate the spin contamination. 

615 

616 Spin contamination is defined as the integral over the 

617 spin density difference, where it is negative (i.e. the 

618 minority spin density is larger than the majority spin density. 

619 """ 

620 

621 if majority_spin == 0: 

622 smaj = 0 

623 smin = 1 

624 else: 

625 smaj = 1 

626 smin = 0 

627 nt_sg, gd = self.get_all_electron_density(atoms) 

628 dt_sg = nt_sg[smin] - nt_sg[smaj] 

629 dt_sg = np.where(dt_sg > 0, dt_sg, 0.0) 

630 return gd.integrate(dt_sg) 

631 

632 def write(self, writer): 

633 writer.write(density=self.gd.collect(self.nt_xG) / Bohr**3, 

634 atomic_density_matrices=pack_atomic_matrices(self.D_asp)) 

635 

636 def read(self, reader): 

637 nt_xG = self.gd.empty(self.ncomponents) 

638 self.gd.distribute(reader.density.density, nt_xG) 

639 nt_xG *= reader.bohr**3 

640 

641 # Read atomic density matrices 

642 natoms = len(self.setups) 

643 atom_partition = AtomPartition(self.gd.comm, np.zeros(natoms, int), 

644 'density-gd') 

645 D_asp = self.setups.empty_atomic_matrix(self.ncomponents, 

646 atom_partition) 

647 self.atom_partition = atom_partition # XXXXXX 

648 spos_ac = np.zeros((natoms, 3)) # XXXX 

649 self.atomdist = self.redistributor.get_atom_distributions(spos_ac) 

650 

651 D_sP = reader.density.atomic_density_matrices 

652 if self.gd.comm.rank == 0: 

653 D_asp.update(unpack_atomic_matrices(D_sP, self.setups, 

654 new=reader.version >= 4, 

655 density=True)) 

656 D_asp.check_consistency() 

657 

658 if self.collinear: 

659 nt_sG = nt_xG 

660 nt_vG = None 

661 else: 

662 nt_sG = nt_xG[:1] 

663 nt_vG = nt_xG[1:] 

664 

665 self.initialize_directly_from_arrays(nt_sG, nt_vG, D_asp) 

666 

667 def initialize_from_other_density(self, dens, kptband_comm): 

668 """Redistribute pseudo density and atomic density matrices. 

669 

670 Collect dens.nt_sG and dens.D_asp to world master and distribute.""" 

671 

672 new_nt_sG = redistribute_array(dens.nt_sG, dens.gd, self.gd, 

673 self.nspins, kptband_comm) 

674 self.atom_partition, self.atomdist = \ 

675 create_atom_partition_and_distibutions(self.gd, self.nspins, 

676 self.setups, 

677 self.redistributor, 

678 kptband_comm) 

679 D_asp = \ 

680 redistribute_atomic_matrices(dens.D_asp, self.gd, self.ncomponents, 

681 self.setups, self.atom_partition, 

682 kptband_comm) 

683 self.initialize_directly_from_arrays(new_nt_sG, None, D_asp) 

684 

685 

686class RealSpaceDensity(Density): 

687 def __init__(self, gd, finegd, nspins, collinear, charge, redistributor, 

688 stencil=3, 

689 background_charge=None): 

690 Density.__init__(self, gd, finegd, nspins, collinear, 

691 charge, redistributor, 

692 background_charge=background_charge) 

693 self.stencil = stencil 

694 

695 self.interpolator = None 

696 

697 def initialize(self, setups, timer, magmom_a, hund): 

698 Density.initialize(self, setups, timer, magmom_a, hund) 

699 

700 # Interpolation function for the density: 

701 self.interpolator = Transformer(self.redistributor.aux_gd, 

702 self.finegd, self.stencil) 

703 

704 spline_aj = [] 

705 for setup in setups: 

706 if setup.nct is None: 

707 spline_aj.append([]) 

708 else: 

709 spline_aj.append([setup.nct]) 

710 self.nct = LFC(self.gd, spline_aj, 

711 integral=[setup.Nct for setup in setups], 

712 forces=True, cut=True) 

713 self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups], 

714 integral=sqrt(4 * pi), forces=True) 

715 

716 def set_positions(self, spos_ac, atom_partition): 

717 Density.set_positions(self, spos_ac, atom_partition) 

718 self.nct_G = self.gd.zeros() 

719 self.nct.add(self.nct_G, 1.0 / self.nspins) 

720 

721 def interpolate_pseudo_density(self, comp_charge=None): 

722 """Interpolate pseudo density to fine grid.""" 

723 if comp_charge is None: 

724 comp_charge, _Q_aL = self.calculate_multipole_moments() 

725 

726 self.nt_sg = self.distribute_and_interpolate(self.nt_sG, self.nt_sg) 

727 

728 # With periodic boundary conditions, the interpolation will 

729 # conserve the number of electrons.