Coverage for gpaw/mixer.py: 93%

1# Copyright (C) 2003 CAMP 

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

3 

4""" 

5See Kresse, Phys. Rev. B 54, 11169 (1996) 

6""" 

7 

8import numpy as np 

9from numpy.fft import fftn, ifftn 

10 

11import gpaw.mpi as mpi 

12from gpaw.utilities.blas import axpy 

13from gpaw.fd_operators import FDOperator 

14from gpaw.utilities.tools import construct_reciprocal 

15from gpaw.new import trace 

16 

17"""About mixing-related classes. 

18 

19(FFT/Broyden)BaseMixer: These classes know how to mix one density 

20array and store history etc. But they do not take care of complexity 

21like spin. 

22 

23(SpinSum/etc.)MixerDriver: These combine one or more BaseMixers to 

24implement a full algorithm. Think of them as stateless (immutable). 

25The user can give an object of these types as input, but they will generally 

26be constructed by a utility function so the interface is nice. 

27 

28The density object always wraps the (X)MixerDriver with a 

29MixerWrapper. The wrapper contains the common code for all mixers so 

30we don't have to implement it multiple times (estimate memory, etc.). 

31 

32In the end, what the user provides is probably a dictionary anyway, and the 

33relevant objects are instantiated automatically.""" 

34 

35 

36class BaseMixer: 

37 name = 'pulay' 

38 

39 """Pulay density mixer.""" 

40 def __init__(self, beta, nmaxold, weight): 

41 """Construct density-mixer object. 

42 

43 Parameters: 

44 

45 beta: float 

46 Mixing parameter between zero and one (one is most 

47 aggressive). 

48 nmaxold: int 

49 Maximum number of old densities. 

50 weight: float 

51 Weight parameter for special metric (for long wave-length 

52 changes). 

53 

54 """ 

55 

56 self.beta = beta 

57 self.nmaxold = nmaxold 

58 self.weight = weight 

59 self.world = None 

60 

61 def initialize_metric(self, gd): 

62 self.gd = gd 

63 

64 if self.weight == 1: 

65 self.metric = None 

66 else: 

67 a = 0.125 * (self.weight + 7) 

68 b = 0.0625 * (self.weight - 1) 

69 c = 0.03125 * (self.weight - 1) 

70 d = 0.015625 * (self.weight - 1) 

71 self.metric = FDOperator([a, 

72 b, b, b, b, b, b, 

73 c, c, c, c, c, c, c, c, c, c, c, c, 

74 d, d, d, d, d, d, d, d], 

75 [(0, 0, 0), # a 

76 (-1, 0, 0), (1, 0, 0), # b 

77 (0, -1, 0), (0, 1, 0), 

78 (0, 0, -1), (0, 0, 1), 

79 (1, 1, 0), (1, 0, 1), (0, 1, 1), # c 

80 (1, -1, 0), (1, 0, -1), (0, 1, -1), 

81 (-1, 1, 0), (-1, 0, 1), (0, -1, 1), 

82 (-1, -1, 0), (-1, 0, -1), (0, -1, -1), 

83 (1, 1, 1), (1, 1, -1), (1, -1, 1), # d 

84 (-1, 1, 1), (1, -1, -1), (-1, -1, 1), 

85 (-1, 1, -1), (-1, -1, -1)], 

86 gd, float).apply 

87 self.mR_sG = gd.empty(4) 

88 

89 def reset(self): 

90 """Reset Density-history. 

91 

92 Called at initialization and after each move of the atoms. 

93 

94 my_nuclei: All nuclei in local domain. 

95 """ 

96 

97 # History for Pulay mixing of densities: 

98 self.nt_isG = [] # Pseudo-electron densities 

99 self.R_isG = [] # Residuals 

100 self.A_ii = np.zeros((0, 0)) 

101 

102 self.D_iasp = [] 

103 self.dD_iasp = [] 

104 

105 def calculate_charge_sloshing(self, R_sG) -> float: 

106 return self.gd.integrate(np.fabs(R_sG)).sum() 

107 

108 def mix_density(self, nt_sG, D_asp, g_ss=None): 

109 nt_isG = self.nt_isG 

110 R_isG = self.R_isG 

111 D_iasp = self.D_iasp 

112 dD_iasp = self.dD_iasp 

113 spin = len(nt_sG) 

114 iold = len(self.nt_isG) 

115 dNt = np.inf 

116 if iold > 0: 

117 if iold > self.nmaxold: 

118 # Throw away too old stuff: 

119 del nt_isG[0] 

120 del R_isG[0] 

121 del D_iasp[0] 

122 del dD_iasp[0] 

123 # for D_p, D_ip, dD_ip in self.D_a: 

124 # del D_ip[0] 

125 # del dD_ip[0] 

126 iold = self.nmaxold 

127 

128 # Calculate new residual (difference between input and 

129 # output density): 

130 R_sG = nt_sG - nt_isG[-1] 

131 dNt = self.calculate_charge_sloshing(R_sG) 

132 R_isG.append(R_sG) 

133 dD_iasp.append([]) 

134 for D_sp, D_isp in zip(D_asp, D_iasp[-1]): 

135 dD_iasp[-1].append(D_sp - D_isp) 

136 

137 if self.metric is None: 

138 mR_sG = R_sG 

139 else: 

140 mR_sG = self.mR_sG[:spin] 

141 for s in range(spin): 

142 self.metric(R_sG[s], mR_sG[s]) 

143 

144 if g_ss is not None: 

145 mR_sG = np.tensordot(g_ss, mR_sG, axes=(1, 0)) 

146 

147 # Update matrix: 

148 A_ii = np.zeros((iold, iold)) 

149 i2 = iold - 1 

150 

151 for i1, R_1sG in enumerate(R_isG): 

152 a = self.gd.comm.sum_scalar( 

153 self.dotprod(R_1sG, mR_sG, dD_iasp[i1], dD_iasp[-1])) 

154 A_ii[i1, i2] = a 

155 A_ii[i2, i1] = a 

156 A_ii[:i2, :i2] = self.A_ii[-i2:, -i2:] 

157 self.A_ii = A_ii 

158 

159 try: 

160 B_ii = np.linalg.inv(A_ii) 

161 alpha_i = B_ii.sum(1) 

162 alpha_i /= alpha_i.sum() 

163 except (ZeroDivisionError, np.linalg.LinAlgError): 

164 alpha_i = np.zeros(iold) 

165 alpha_i[-1] = 1.0 

166 

167 if self.world: 

168 self.world.broadcast(alpha_i, 0) 

169 

170 # Calculate new input density: 

171 nt_sG[:] = 0.0 

172 # for D_p, D_ip, dD_ip in self.D_a: 

173 for D in D_asp: 

174 D[:] = 0.0 

175 beta = self.beta 

176 for i, alpha in enumerate(alpha_i): 

177 axpy(alpha, nt_isG[i], nt_sG) 

178 axpy(alpha * beta, R_isG[i], nt_sG) 

179 

180 for D_sp, D_isp, dD_isp in zip(D_asp, D_iasp[i], 

181 dD_iasp[i]): 

182 axpy(alpha, D_isp, D_sp) 

183 axpy(alpha * beta, dD_isp, D_sp) 

184 

185 # Store new input density (and new atomic density matrices): 

186 nt_isG.append(nt_sG.copy()) 

187 D_iasp.append([]) 

188 for D_sp in D_asp: 

189 D_iasp[-1].append(D_sp.copy()) 

190 return dNt 

191 

192 # may presently be overridden by passing argument in constructor 

193 def dotprod(self, R1_G, R2_G, dD1_ap, dD2_ap): 

194 return np.vdot(R1_G, R2_G).real 

195 

196 def estimate_memory(self, mem, gd): 

197 gridbytes = gd.bytecount() 

198 mem.subnode('nt_iG, R_iG', 2 * self.nmaxold * gridbytes) 

199 

200 def __repr__(self): 

201 classname = self.__class__.__name__ 

202 template = '%s(beta=%f, nmaxold=%d, weight=%f)' 

203 string = template % (classname, self.beta, self.nmaxold, self.weight) 

204 return string 

205 

206 

207class ExperimentalDotProd: 

208 def __init__(self, calc): 

209 self.calc = calc 

210 

211 def __call__(self, R1_G, R2_G, dD1_ap, dD2_ap): 

212 prod = np.vdot(R1_G, R2_G).real 

213 setups = self.calc.wfs.setups 

214 # okay, this is a bit nasty because it depends on dD1_ap 

215 # and its friend having come from D_asp.values() and the dictionaries 

216 # not having been modified. This is probably true... for now. 

217 avalues = self.calc.density.D_asp.keys() 

218 for a, dD1_p, dD2_p in zip(avalues, dD1_ap, dD2_ap): 

219 I4_pp = setups[a].four_phi_integrals() 

220 dD4_pp = np.outer(dD1_p, dD2_p) # not sure if corresponds quite 

221 prod += (I4_pp * dD4_pp).sum() 

222 return prod 

223 

224 

225class ReciprocalMetric: 

226 def __init__(self, weight, k2_Q): 

227 self.k2_Q = k2_Q 

228 k2_min = np.min(self.k2_Q) 

229 self.q1 = (weight - 1) * k2_min 

230 

231 def __call__(self, R_Q, mR_Q): 

232 mR_Q[:] = R_Q * (1.0 + self.q1 / self.k2_Q) 

233 

234 

235class FFTBaseMixer(BaseMixer): 

236 name = 'fft' 

237 

238 """Mix the density in Fourier space""" 

239 def __init__(self, beta, nmaxold, weight): 

240 BaseMixer.__init__(self, beta, nmaxold, weight) 

241 self.gd1 = None 

242 

243 def initialize_metric(self, gd): 

244 self.gd = gd 

245 

246 if gd.comm.rank == 0: 

247 self.gd1 = gd.new_descriptor(comm=mpi.serial_comm) 

248 k2_Q, _ = construct_reciprocal(self.gd1) 

249 self.metric = ReciprocalMetric(self.weight, k2_Q) 

250 self.mR_sG = self.gd1.empty(2, dtype=complex) 

251 else: 

252 self.metric = lambda R_Q, mR_Q: None 

253 self.mR_sG = np.empty((2, 0, 0, 0), dtype=complex) 

254 

255 def calculate_charge_sloshing(self, R_sQ): 

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

257 assert R_sQ.ndim == 4 # and len(R_sQ) == 1 

258 cs = sum(self.gd1.integrate(np.fabs(ifftn(R_Q).real)) 

259 for R_Q in R_sQ) 

260 else: 

261 cs = 0.0 

262 return self.gd.comm.sum_scalar(cs) 

263 

264 def mix_density(self, nt_sR, D_asp, g_ss=None): 

265 # Transform real-space density to Fourier space 

266 nt1_sR = [self.gd.collect(nt_R) for nt_R in nt_sR] 

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

268 nt1_sG = np.ascontiguousarray([fftn(nt_R) for nt_R in nt1_sR]) 

269 else: 

270 nt1_sG = np.empty((len(nt_sR), 0, 0, 0), dtype=complex) 

271 

272 dNt = BaseMixer.mix_density(self, nt1_sG, D_asp) 

273 

274 # Return density in real space 

275 for nt_G, nt_R in zip(nt1_sG, nt_sR): 

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

277 nt1_R = ifftn(nt_G).real 

278 else: 

279 nt1_R = None 

280 self.gd.distribute(nt1_R, nt_R) 

281 

282 return dNt 

283 

284 

285class BroydenBaseMixer: 

286 name = 'broyden' 

287 

288 def __init__(self, beta, nmaxold, weight): 

289 self.verbose = False 

290 self.beta = beta 

291 self.nmaxold = nmaxold 

292 self.weight = 1.0 # XXX discards argument 

293 

294 def initialize_metric(self, gd): 

295 self.gd = gd 

296 

297 def reset(self): 

298 self.step = 0 

299 # self.d_nt_G = [] 

300 # self.d_D_ap = [] 

301 

302 self.R_iG = [] 

303 self.dD_iap = [] 

304 

305 self.nt_iG = [] 

306 self.D_iap = [] 

307 self.c_G = [] 

308 self.v_G = [] 

309 self.u_G = [] 

310 self.u_D = [] 

311 

312 def mix_density(self, nt_sG, D_asp): 

313 nt_G = nt_sG[0] 

314 D_ap = [D_sp[0] for D_sp in D_asp] 

315 dNt = np.inf 

316 if self.step > 2: 

317 del self.R_iG[0] 

318 for d_Dp in self.dD_iap: 

319 del d_Dp[0] 

320 if self.step > 0: 

321 self.R_iG.append(nt_G - self.nt_iG[-1]) 

322 for d_Dp, D_p, D_ip in zip(self.dD_iap, D_ap, self.D_iap): 

323 d_Dp.append(D_p - D_ip[-1]) 

324 fmin_G = self.gd.integrate(self.R_iG[-1] * self.R_iG[-1]) 

325 dNt = self.gd.integrate(np.fabs(self.R_iG[-1])) 

326 if self.verbose: 

327 print('Mixer: broyden: fmin_G = %f fmin_D = %f' % fmin_G) 

328 if self.step == 0: 

329 self.eta_G = np.empty(nt_G.shape) 

330 self.eta_D = [] 

331 for D_p in D_ap: 

332 self.eta_D.append(0) 

333 self.u_D.append([]) 

334 self.D_iap.append([]) 

335 self.dD_iap.append([]) 

336 else: 

337 if self.step >= 2: 

338 del self.c_G[:] 

339 if len(self.v_G) >= self.nmaxold: 

340 del self.u_G[0] 

341 del self.v_G[0] 

342 for u_D in self.u_D: 

343 del u_D[0] 

344 temp_nt_G = self.R_iG[1] - self.R_iG[0] 

345 self.v_G.append(temp_nt_G / self.gd.integrate(temp_nt_G * 

346 temp_nt_G)) 

347 if len(self.v_G) < self.nmaxold: 

348 nstep = self.step - 1 

349 else: 

350 nstep = self.nmaxold 

351 for i in range(nstep): 

352 self.c_G.append(self.gd.integrate(self.v_G[i] * 

353 self.R_iG[1])) 

354 self.u_G.append(self.beta * temp_nt_G + self.nt_iG[1] - 

355 self.nt_iG[0]) 

356 for d_Dp, u_D, D_ip in zip(self.dD_iap, self.u_D, self.D_iap): 

357 temp_D_ap = d_Dp[1] - d_Dp[0] 

358 u_D.append(self.beta * temp_D_ap + D_ip[1] - D_ip[0]) 

359 usize = len(self.u_G) 

360 for i in range(usize - 1): 

361 a_G = self.gd.integrate(self.v_G[i] * temp_nt_G) 

362 axpy(-a_G, self.u_G[i], self.u_G[usize - 1]) 

363 for u_D in self.u_D: 

364 axpy(-a_G, u_D[i], u_D[usize - 1]) 

365 self.eta_G = self.beta * self.R_iG[-1] 

366 for i, d_Dp in enumerate(self.dD_iap): 

367 self.eta_D[i] = self.beta * d_Dp[-1] 

368 usize = len(self.u_G) 

369 for i in range(usize): 

370 axpy(-self.c_G[i], self.u_G[i], self.eta_G) 

371 for eta_D, u_D in zip(self.eta_D, self.u_D): 

372 axpy(-self.c_G[i], u_D[i], eta_D) 

373 axpy(-1.0, self.R_iG[-1], nt_G) 

374 axpy(1.0, self.eta_G, nt_G) 

375 for D_p, d_Dp, eta_D in zip(D_ap, self.dD_iap, self.eta_D): 

376 axpy(-1.0, d_Dp[-1], D_p) 

377 axpy(1.0, eta_D, D_p) 

378 if self.step >= 2: 

379 del self.nt_iG[0] 

380 for D_ip in self.D_iap: 

381 del D_ip[0] 

382 self.nt_iG.append(np.copy(nt_G)) 

383 for D_ip, D_p in zip(self.D_iap, D_ap): 

384 D_ip.append(np.copy(D_p)) 

385 self.step += 1 

386 return dNt 

387 

388 

389class DummyMixer: 

390 """Dummy mixer for TDDFT, i.e., it does not mix.""" 

391 name = 'dummy' 

392 beta = 1.0 

393 nmaxold = 1 

394 weight = 1 

395 

396 def __init__(self, *args, **kwargs): 

397 return 

398 

399 def mix(self, basemixers, nt_sG, D_asp): 

400 return 0.0 

401 

402 def get_basemixers(self, nspins): 

403 return [] 

404 

405 def todict(self): 

406 return {'name': 'dummy'} 

407 

408 

409class NotMixingMixer: 

410 name = 'no-mixing' 

411 

412 def __init__(self, beta, nmaxold, weight): 

413 """Construct density-mixer object. 

414 Parameters: they are ignored for this mixer 

415 """ 

416 

417 # whatever parameters you give it doesn't do anything with them 

418 self.beta = 0 

419 self.nmaxold = 0 

420 self.weight = 0 

421 

422 def initialize_metric(self, gd): 

423 self.gd = gd 

424 self.metric = None 

425 

426 def reset(self): 

427 """Reset Density-history. 

428 

429 Called at initialization and after each move of the atoms. 

430 

431 my_nuclei: All nuclei in local domain. 

432 """ 

433 

434 # Previous density: 

435 self.nt_isG = [] # Pseudo-electron densities 

436 

437 def calculate_charge_sloshing(self, R_sG): 

438 return self.gd.integrate(np.fabs(R_sG)).sum() 

439 

440 def mix_density(self, nt_sG, D_asp): 

441 iold = len(self.nt_isG) 

442 

443 dNt = np.inf 

444 if iold > 0: 

445 # Calculate new residual (difference between input and 

446 # output density): 

447 dNt = self.calculate_charge_sloshing(nt_sG - self.nt_isG[-1]) 

448 # Store new input density: 

449 self.nt_isG = [nt_sG.copy()] 

450 

451 return dNt 

452 

453 # may presently be overridden by passing argument in constructor 

454 def dotprod(self, R1_G, R2_G, dD1_ap, dD2_ap): 

455 pass 

456 

457 def estimate_memory(self, mem, gd): 

458 gridbytes = gd.bytecount() 

459 mem.subnode('nt_iG, R_iG', 2 * self.nmaxold * gridbytes) 

460 

461 def __repr__(self): 

462 string = 'no mixing of density' 

463 return string 

464 

465 

466class SeparateSpinMixerDriver: 

467 name = 'separate' 

468 

469 def __init__(self, basemixerclass, beta, nmaxold, weight): 

470 self.basemixerclass = basemixerclass 

471 

472 self.beta = beta 

473 self.nmaxold = nmaxold 

474 self.weight = weight 

475 

476 def get_basemixers(self, nspins): 

477 return [self.basemixerclass(self.beta, self.nmaxold, self.weight) 

478 for _ in range(nspins)] 

479 

480 def mix(self, basemixers, nt_sG, D_asp): 

481 """Mix pseudo electron densities.""" 

482 D_asp = D_asp.values() 

483 dNt = 0.0 

484 for s, (nt_G, basemixer) in enumerate(zip(nt_sG, basemixers)): 

485 D_a1p = [D_sp[s:s + 1] for D_sp in D_asp] 

486 nt_1G = nt_G[np.newaxis] 

487 dNt += basemixer.mix_density(nt_1G, D_a1p) 

488 return dNt 

489 

490 

491class SpinSumMixerDriver: 

492 name = 'sum' 

493 mix_atomic_density_matrices = False 

494 

495 def __init__(self, basemixerclass, beta, nmaxold, weight): 

496 self.basemixerclass = basemixerclass 

497 

498 self.beta = beta 

499 self.nmaxold = nmaxold 

500 self.weight = weight 

501 

502 def get_basemixers(self, nspins): 

503 if nspins == 1: 

504 raise ValueError('Spin sum mixer expects 2 or 4 components') 

505 return [self.basemixerclass(self.beta, self.nmaxold, self.weight)] 

506 

507 def mix(self, basemixers, nt_sG, D_asp): 

508 assert len(basemixers) == 1 

509 basemixer = basemixers[0] 

510 D_asp = D_asp.values() 

511 

512 collinear = len(nt_sG) == 2 

513 

514 # Mix density 

515 if collinear: 

516 nt_1G = nt_sG.sum(0)[np.newaxis] 

517 else: 

518 nt_1G = nt_sG[:1] 

519 

520 if self.mix_atomic_density_matrices: 

521 if collinear: 

522 D_a1p = [D_sp[:1] + D_sp[1:] for D_sp in D_asp] 

523 else: 

524 D_a1p = [D_sp[:1] for D_sp in D_asp] 

525 dNt = basemixer.mix_density(nt_1G, D_a1p) 

526 if collinear: 

527 dD_ap = [D_sp[0] - D_sp[1] for D_sp in D_asp] 

528 for D_sp, D_1p, dD_p in zip(D_asp, D_a1p, dD_ap): 

529 D_sp[0] = 0.5 * (D_1p[0] + dD_p) 

530 D_sp[1] = 0.5 * (D_1p[0] - dD_p) 

531 else: 

532 dNt = basemixer.mix_density(nt_1G, D_asp) 

533 

534 if collinear: 

535 dnt_G = nt_sG[0] - nt_sG[1] 

536 # Only new magnetization for spin density 

537 # dD_ap = [D_sp[0] - D_sp[1] for D_sp in D_asp] 

538 

539 # Construct new spin up/down densities 

540 nt_sG[0] = 0.5 * (nt_1G[0] + dnt_G) 

541 nt_sG[1] = 0.5 * (nt_1G[0] - dnt_G) 

542 

543 return dNt 

544 

545 

546class SpinSumMixerDriver2(SpinSumMixerDriver): 

547 name = 'sum2' 

548 mix_atomic_density_matrices = True 

549 

550 

551class SpinDifferenceMixerDriver: 

552 name = 'difference' 

553 

554 def __init__(self, basemixerclass, beta, nmaxold, weight, 

555 beta_m=0.7, nmaxold_m=2, weight_m=10.0): 

556 self.basemixerclass = basemixerclass 

557 self.beta = beta 

558 self.nmaxold = nmaxold 

559 self.weight = weight 

560 self.beta_m = beta_m 

561 self.nmaxold_m = nmaxold_m 

562 self.weight_m = weight_m 

563 

564 def get_basemixers(self, nspins): 

565 if nspins == 1: 

566 raise ValueError('Spin difference mixer expects 2 or 4 components') 

567 basemixer = self.basemixerclass(self.beta, self.nmaxold, self.weight) 

568 if nspins == 2: 

569 basemixer_m = self.basemixerclass(self.beta_m, self.nmaxold_m, 

570 self.weight_m) 

571 return basemixer, basemixer_m 

572 else: 

573 basemixer_x = self.basemixerclass(self.beta_m, self.nmaxold_m, 

574 self.weight_m) 

575 basemixer_y = self.basemixerclass(self.beta_m, self.nmaxold_m, 

576 self.weight_m) 

577 basemixer_z = self.basemixerclass(self.beta_m, self.nmaxold_m, 

578 self.weight_m) 

579 return basemixer, basemixer_x, basemixer_y, basemixer_z 

580 

581 def mix(self, basemixers, nt_sG, D_asp): 

582 D_asp = D_asp.values() 

583 

584 if len(nt_sG) == 2: 

585 basemixer, basemixer_m = basemixers 

586 else: 

587 assert len(nt_sG) == 4 

588 basemixer, basemixer_x, basemixer_y, basemixer_z = basemixers 

589 

590 if len(nt_sG) == 2: 

591 # Mix density 

592 nt_1G = nt_sG.sum(0)[np.newaxis] 

593 D_a1p = [D_sp[:1] + D_sp[1:] for D_sp in D_asp] 

594 dNt = basemixer.mix_density(nt_1G, D_a1p) 

595 

596 # Mix magnetization 

597 dnt_1G = nt_sG[:1] - nt_sG[1:] 

598 dD_a1p = [D_sp[:1] - D_sp[1:] for D_sp in D_asp] 

599 basemixer_m.mix_density(dnt_1G, dD_a1p) 

600 # (The latter is not counted in dNt) 

601 

602 # Construct new spin up/down densities 

603 nt_sG[:1] = 0.5 * (nt_1G + dnt_1G) 

604 nt_sG[1:] = 0.5 * (nt_1G - dnt_1G) 

605 for D_sp, D_1p, dD_1p in zip(D_asp, D_a1p, dD_a1p): 

606 D_sp[:1] = 0.5 * (D_1p + dD_1p) 

607 D_sp[1:] = 0.5 * (D_1p - dD_1p) 

608 else: 

609 # Mix density 

610 nt_1G = nt_sG[:1] 

611 D_a1p = [D_sp[:1] for D_sp in D_asp] 

612 dNt = basemixer.mix_density(nt_1G, D_a1p) 

613 

614 # Mix magnetization 

615 Dx_a1p = [D_sp[1:2] for D_sp in D_asp] 

616 Dy_a1p = [D_sp[2:3] for D_sp in D_asp] 

617 Dz_a1p = [D_sp[3:4] for D_sp in D_asp] 

618 

619 basemixer_x.mix_density(nt_sG[1:2], Dx_a1p) 

620 basemixer_y.mix_density(nt_sG[2:3], Dy_a1p) 

621 basemixer_z.mix_density(nt_sG[3:4], Dz_a1p) 

622 return dNt 

623 

624 

625class FullSpinMixerDriver: 

626 name = 'fullspin' 

627 

628 def __init__(self, basemixerclass, beta, nmaxold, weight, g=None): 

629 self.basemixerclass = basemixerclass 

630 self.beta = beta 

631 self.nmaxold = nmaxold 

632 self.weight = weight 

633 self.g_ss = g 

634 

635 def get_basemixers(self, nspins): 

636 if nspins == 1: 

637 raise ValueError('Full-spin mixer expects 2 or 4 spin channels') 

638 

639 basemixer = self.basemixerclass(self.beta, self.nmaxold, self.weight) 

640 return [basemixer] 

641 

642 def mix(self, basemixers, nt_sG, D_asp): 

643 D_asp = D_asp.values() 

644 basemixer = basemixers[0] 

645 if self.g_ss is None: 

646 self.g_ss = np.identity(len(nt_sG)) 

647 

648 dNt = basemixer.mix_density(nt_sG, D_asp, self.g_ss) 

649 

650 return dNt 

651 

652 

653# Dictionaries to get mixers by name: 

654_backends = {} 

655_methods = {} 

656for cls in [FFTBaseMixer, BroydenBaseMixer, BaseMixer, NotMixingMixer]: 

657 _backends[cls.name] = cls # type:ignore 

658for dcls in [SeparateSpinMixerDriver, SpinSumMixerDriver, 

659 FullSpinMixerDriver, SpinSumMixerDriver2, 

660 SpinDifferenceMixerDriver, DummyMixer]: 

661 _methods[dcls.name] = dcls # type:ignore 

662 

663 

664# This function is used by Density to decide mixer parameters 

665# that the user did not explicitly provide, i.e., it fills out 

666# everything that is missing and returns a mixer "driver". 

667def get_mixer_from_keywords(pbc, nspins, **mixerkwargs): 

668 if mixerkwargs.get('name') == 'dummy': 

669 return DummyMixer() 

670 

671 if mixerkwargs.get('backend') == 'no-mixing': 

672 mixerkwargs['beta'] = 0 

673 mixerkwargs['nmaxold'] = 0 

674 mixerkwargs['weight'] = 0 

675 

676 if nspins == 1: 

677 mixerkwargs['method'] = SeparateSpinMixerDriver 

678 

679 # The plan is to first establish a kwargs dictionary with all the 

680 # defaults, then we update it with values from the user. 

681 kwargs = {'backend': BaseMixer} 

682 

683 if np.any(pbc): # Works on array or boolean 

684 kwargs.update(beta=0.05, history=5, weight=50.0) 

685 else: 

686 kwargs.update(beta=0.25, history=3, weight=1.0) 

687 

688 if nspins == 1: 

689 kwargs['method'] = SeparateSpinMixerDriver 

690 else: 

691 kwargs['method'] = SpinDifferenceMixerDriver 

692 

693 # Clean up mixerkwargs (compatibility) 

694 if 'nmaxold' in mixerkwargs: 

695 assert 'history' not in mixerkwargs 

696 mixerkwargs['history'] = mixerkwargs.pop('nmaxold') 

697 

698 # Now the user override: 

699 for key in kwargs: 

700 # Clean any 'None' values out as if they had never been passed: 

701 val = mixerkwargs.pop(key, None) 

702 if val is not None: 

703 kwargs[key] = val 

704 

705 # Resolve keyword strings (like 'fft') into classes (like FFTBaseMixer): 

706 driver = _methods.get(kwargs['method'], kwargs['method']) 

707 baseclass = _backends.get(kwargs['backend'], kwargs['backend']) 

708 

709 # We forward any remaining mixer kwargs to the actual mixer object. 

710 # Any user defined variables that do not really exist will cause an error. 

711 mixer = driver(baseclass, beta=kwargs['beta'], 

712 nmaxold=kwargs['history'], weight=kwargs['weight'], 

713 **mixerkwargs) 

714 return mixer 

715 

716