Coverage for gpaw/fdtd/poisson_fdtd.py: 90%

1"""Part of the module for electrodynamic simulations.""" 

2 

3import numpy as np 

4from ase import Atoms 

5from ase.io.trajectory import read_atoms 

6from ase.units import Bohr 

7 

8import gpaw.mpi as mpi 

9from gpaw import GPAW 

10from gpaw.fdtd.polarizable_material import PolarizableMaterial 

11from gpaw.fdtd.potential_couplers import (RefinerPotentialCoupler, 

12 MultipolesPotentialCoupler) 

13from gpaw.grid_descriptor import GridDescriptor 

14from gpaw.mpi import world, serial_comm 

15from gpaw.poisson import PoissonSolver 

16from gpaw.poisson_moment import MomentCorrectionPoissonSolver 

17from gpaw.tddft import TDDFT, DipoleMomentWriter, RestartFileWriter 

18from gpaw.transformers import Transformer 

19from gpaw.utilities import h2gpts 

20 

21# QSFDTD is a wrapper class to make calculations 

22# easier, something like this: 

23# qsfdtd = QSFDTD(cl, at, c, h) 

24# energy = qsfdtd.ground_state('gs.gpw') 

25# qsfdtd.time_propagation('gs.gpw', kick) 

26# photoabsorption_spectrum('dm.dat', 'spec.dat') 

27 

28 

29class QSFDTD: 

30 def __init__(self, 

31 classical_material, 

32 atoms, 

33 cells, 

34 spacings, 

35 communicator=serial_comm, 

36 remove_moments=(1, 1)): 

37 

38 self.td_calc = None 

39 

40 # Define classical cell in one of these ways: 

41 # 1. [cx, cy, cz] -> vacuum for the quantum grid = 4.0 

42 # 2. ([cx, cy, cz]) -> vacuum = 4.0 

43 # 3. ([cx, cy, cz], vacuum) 

44 # 4. ([cx, cy, cz], ([p1x, p1y, p1z], [p2x, p2y, p2z])) 

45 # where p1 and p2 are corners of the quantum grid 

46 if len(cells) == 1: # (cell) 

47 assert (len(cells[0]) == 3) 

48 cell = np.array(cells[0]) 

49 vacuum = 4.0 

50 elif len(cells) == 3: # [cell.x, cell.y, cell.z] 

51 cell = np.array(cells) 

52 vacuum = 4.0 

53 elif len(cells) == 2: # cell + vacuum/corners 

54 cell = np.array(cells[0]) 

55 if np.array(cells[1]).size == 1: 

56 vacuum = cells[1] 

57 corners = None 

58 else: 

59 vacuum = None 

60 corners = cells[1] 

61 else: 

62 raise Exception('QSFDTD: cells defined incorrectly') 

63 

64 # Define spacings in in one of these ways: 

65 # 1. (clh, qmh) 

66 # 2. clh -> qmh=clh/4 

67 if np.array(spacings).size == 1: # clh 

68 cl_spacing = spacings 

69 qm_spacing = 0.25 * cl_spacing 

70 elif len(spacings) == 2: # (clh, qmh) 

71 cl_spacing = spacings[0] 

72 qm_spacing = spacings[1] 

73 else: 

74 raise Exception('QSFDTD: spacings defined incorrectly') 

75 

76 self.poissonsolver = FDTDPoissonSolver( 

77 classical_material=classical_material, 

78 cl_spacing=cl_spacing, 

79 qm_spacing=qm_spacing, 

80 remove_moments=remove_moments, 

81 communicator=communicator, 

82 cell=cell) 

83 self.poissonsolver.set_calculation_mode('iterate') 

84 

85 # Quantum system 

86 if atoms is not None: 

87 assert (len(cells) == 2) 

88 assert (len(spacings) == 2) 

89 assert (len(remove_moments) == 2) 

90 self.atoms = atoms 

91 self.atoms.set_cell(cell) 

92 if vacuum is not None: # vacuum 

93 self.atoms, self.qm_spacing, self.gpts = \ 

94 self.poissonsolver.cut_cell(self.atoms, vacuum=vacuum) 

95 else: # corners 

96 self.atoms, self.qm_spacing, self.gpts = \ 

97 self.poissonsolver.cut_cell(self.atoms, corners=corners) 

98 else: # Dummy quantum system 

99 self.atoms = Atoms('H', [0.5 * cell], cell=cell) 

100 if vacuum is not None: # vacuum 

101 self.atoms, self.qm_spacing, self.gpts = \ 

102 self.poissonsolver.cut_cell(self.atoms, vacuum=vacuum) 

103 else: # corners 

104 self.atoms, self.qm_spacing, self.gpts = \ 

105 self.poissonsolver.cut_cell(self.atoms, corners=corners) 

106 del self.atoms[:] 

107 

108 def ground_state(self, filename, **kwargs): 

109 # GPAW calculator for the ground state 

110 self.gs_calc = GPAW(gpts=self.gpts, 

111 poissonsolver=self.poissonsolver, 

112 **kwargs) 

113 self.atoms.calc = self.gs_calc 

114 self.energy = self.atoms.get_potential_energy() 

115 self.write(filename, mode='all') 

116 

117 def write(self, filename, **kwargs): 

118 if self.td_calc is None: 

119 self.gs_calc.write(filename, **kwargs) 

120 else: 

121 self.td_calc.write(filename, **kwargs) 

122 

123 def time_propagation(self, 

124 filename, 

125 kick_strength, 

126 time_step, 

127 iterations, 

128 dipole_moment_file=None, 

129 restart_file=None, 

130 dump_interval=100, 

131 **kwargs): 

132 self.td_calc = TDDFT(filename, **kwargs) 

133 if dipole_moment_file is not None: 

134 DipoleMomentWriter(self.td_calc, dipole_moment_file) 

135 if restart_file is not None: 

136 RestartFileWriter(self.td_calc, restart_file, dump_interval) 

137 if kick_strength is not None: 

138 self.td_calc.absorption_kick(kick_strength) 

139 self.td_calc.propagate(time_step, iterations) 

140 if restart_file is not None: 

141 self.td_calc.write(restart_file, 'all') 

142 

143 

144# This helps in telling the classical quantities from the quantum ones 

145class PoissonOrganizer: 

146 def __init__(self): 

147 self.poisson_solver = None 

148 self.gd = None 

149 self.density = None 

150 self.cell = None 

151 self.spacing_def = None 

152 self.spacing = None 

153 

154 

155# For mixing the potential 

156class SimpleMixer(): 

157 def __init__(self, alpha, data): 

158 self.alpha = alpha 

159 self.data = np.copy(data) 

160 

161 def mix(self, data): 

162 self.data = self.alpha * data + (1.0 - self.alpha) * self.data 

163 return np.copy(self.data) 

164 

165 

166# Contains one PoissonSolver for the classical 

167# and one for the quantum subsystem 

168class FDTDPoissonSolver: 

169 def __init__(self, 

170 nn=3, 

171 relax='J', 

172 eps=1e-24, 

173 classical_material=None, 

174 cell=None, 

175 qm_spacing=0.30, 

176 cl_spacing=1.20, 

177 remove_moments=(1, 1), 

178 potential_coupler='Refiner', 

179 communicator=serial_comm): 

180 

181 self.rank = mpi.rank 

182 

183 self.messages = [] 

184 

185 assert (potential_coupler in ['Multipoles', 'Refiner']) 

186 self.potential_coupling_scheme = potential_coupler 

187 

188 if classical_material is None: 

189 self.classical_material = PolarizableMaterial() 

190 else: 

191 self.classical_material = classical_material 

192 

193 self.set_calculation_mode('solve') 

194 

195 self.has_subsystems = False 

196 self.remove_moment_cl = remove_moments[0] 

197 self.remove_moment_qm = remove_moments[1] 

198 self.time = 0.0 

199 self.time_step = 0.0 

200 self.kick = np.array([0.0, 0.0, 0.0], dtype=float) 

201 self.maxiter = 2000 

202 self.eps = eps 

203 self.relax = relax 

204 self.nn = nn 

205 

206 # Only handle the quantities via self.qm or self.cl 

207 self.cl = PoissonOrganizer() 

208 self.cl.spacing_def = cl_spacing * np.ones(3) / Bohr 

209 self.cl.extrapolated_qm_phi = None 

210 self.cl.dcomm = communicator 

211 self.cl.dparsize = None 

212 self.qm = PoissonOrganizer() 

213 self.qm.spacing_def = qm_spacing * np.ones(3) / Bohr 

214 self.qm.cell = np.array(cell) / Bohr 

215 

216 # Classical spacing and simulation cell 

217 _cell = np.array(cell) / Bohr 

218 self.cl.spacing = self.cl.spacing_def 

219 if np.size(_cell) == 3: 

220 self.cl.cell = np.diag(_cell) 

221 else: 

222 self.cl.cell = _cell 

223 

224 # Generate classical grid descriptor 

225 self.initialize_clgd() 

226 

227 def todict(self): 

228 dct = dict( 

229 name='fdtd', 

230 nn=self.nn, 

231 relax=self.relax, 

232 eps=self.eps, 

233 # classical_material missing 

234 cell=self.cl.cell * Bohr, 

235 qm_spacing=self.qm.spacing_def[0] * Bohr, 

236 cl_spacing=self.cl.spacing_def[0] * Bohr, 

237 remove_moments=(self.remove_moment_cl, self.remove_moment_qm), 

238 potential_coupler=self.potential_coupling_scheme) 

239 

240 return dct 

241 

242 def get_description(self): 

243 return 'FDTD+TDDFT' 

244 

245 # Generated classical GridDescriptors, after spacing and cell are known 

246 def initialize_clgd(self): 

247 N_c = h2gpts(self.cl.spacing, self.cl.cell, 4) 

248 self.cl.spacing = np.diag(self.cl.cell) / N_c 

249 self.cl.gd = GridDescriptor(N_c, self.cl.cell, False, self.cl.dcomm, 

250 self.cl.dparsize) 

251 self.cl.gd_global = GridDescriptor(N_c, self.cl.cell, False, 

252 serial_comm, None) 

253 self.cl.extrapolated_qm_phi = self.cl.gd.empty() 

254 

255 def estimate_memory(self, mem): 

256 # self.cl.poisson_solver.estimate_memory(mem) 

257 # print(self.qm.poisson_solver.estimate_memory.__code__.co_varnames) 

258 # WTF? How can this shabby method suddenly be unbound? 

259 # It needs both self and mem. 

260 # Ferchrissakes! 

261 # self.qm.poisson_solver.estimate_memory(mem=mem) 

262 pass 

263 

264 # Return the TDDFT stencil by default 

265 def get_stencil(self, mode='qm'): 

266 if mode == 'qm': 

267 return self.qm.poisson_solver.get_stencil() 

268 else: 

269 return self.cl.poisson_solver.get_stencil() 

270 

271 # Initialize both PoissonSolvers 

272 # def initialize(self): 

273 # self.qm.poisson_solver._init() 

274 # self.cl.poisson_solver._init() 

275 

276 def set_grid_descriptor(self, qmgd): 

277 if not self.has_subsystems: 

278 return 

279 

280 self.qm.gd = qmgd 

281 

282 # Create quantum Poisson solver 

283 poisson_qm_kwargs = dict(name='fd', nn=self.nn, 

284 eps=self.eps, relax=self.relax) 

285 self.qm.poisson_solver = MomentCorrectionPoissonSolver( 

286 poissonsolver=PoissonSolver(**poisson_qm_kwargs), 

287 moment_corrections=self.remove_moment_qm) 

288 self.qm.poisson_solver.set_grid_descriptor(self.qm.gd) 

289 # self.qm.poisson_solver.initialize() 

290 self.qm.phi = self.qm.gd.zeros() 

291 self.qm.rho = self.qm.gd.zeros() 

292 

293 # Set quantum grid descriptor 

294 self.qm.poisson_solver.set_grid_descriptor(qmgd) 

295 

296 # Create classical PoissonSolver 

297 poisson_cl_kwargs = dict(name='fd', nn=self.nn, 

298 eps=self.eps, relax=self.relax) 

299 self.cl.poisson_solver = MomentCorrectionPoissonSolver( 

300 poissonsolver=PoissonSolver(**poisson_cl_kwargs), 

301 moment_corrections=self.remove_moment_cl) 

302 self.cl.poisson_solver.set_grid_descriptor(self.cl.gd) 

303 # self.cl.poisson_solver.initialize() 

304 

305 # Initialize classical material, 

306 # its Poisson solver was generated already 

307 self.cl.poisson_solver.set_grid_descriptor(self.cl.gd) 

308 self.classical_material.initialize(self.cl.gd) 

309 self.cl.extrapolated_qm_phi = self.cl.gd.zeros() 

310 self.cl.phi = self.cl.gd.zeros() 

311 self.cl.extrapolated_qm_phi = self.cl.gd.empty() 

312 

313 msg = self.messages.append 

314 msg('\nFDTDPoissonSolver/grid descriptors and coupler:') 

315 msg(' Domain parallelization with %i processes.' % 

316 self.cl.gd.comm.size) 

317 if self.cl.gd.comm == serial_comm: 

318 msg(' Communicator for domain parallelization: serial_comm') 

319 elif self.cl.gd.comm == world: 

320 msg(' Communicator for domain parallelization: world') 

321 elif self.cl.gd.comm == self.qm.gd.comm: 

322 msg(' Communicator for domain parallelization: dft_domain_comm') 

323 else: 

324 msg(' Communicator for domain parallelization: %s' 

325 % self.cl.gd.comm) 

326 

327 # Initialize potential coupler 

328 if self.potential_coupling_scheme == 'Multipoles': 

329 msg('Classical-quantum coupling by multipole expansion ' + 

330 'with maxL: %i' % (self.remove_moment_qm)) 

331 self.potential_coupler = MultipolesPotentialCoupler( 

332 qm=self.qm, 

333 cl=self.cl, 

334 index_offset_1=self.shift_indices_1, 

335 index_offset_2=self.shift_indices_2, 

336 extended_index_offset_1=self.extended_shift_indices_1, 

337 extended_index_offset_2=self.extended_shift_indices_2, 

338 extended_delta_index=self.extended_deltaIndex, 

339 num_refinements=self.num_refinements, 

340 remove_moment_qm=self.remove_moment_qm, 

341 remove_moment_cl=self.remove_moment_cl, 

342 rank=self.rank) 

343 else: 

344 msg('Classical-quantum coupling by coarsening/refining') 

345 self.potential_coupler = RefinerPotentialCoupler( 

346 qm=self.qm, 

347 cl=self.cl, 

348 index_offset_1=self.shift_indices_1, 

349 index_offset_2=self.shift_indices_2, 

350 extended_index_offset_1=self.extended_shift_indices_1, 

351 extended_index_offset_2=self.extended_shift_indices_2, 

352 extended_delta_index=self.extended_deltaIndex, 

353 num_refinements=self.num_refinements, 

354 remove_moment_qm=self.remove_moment_qm, 

355 remove_moment_cl=self.remove_moment_cl, 

356 rank=self.rank) 

357 

358 self.phi_tot_clgd = self.cl.gd.empty() 

359 self.phi_tot_qmgd = self.qm.gd.empty() 

360 

361 def cut_cell(self, 

362 atoms_in, 

363 vacuum=5.0, 

364 corners=None, 

365 create_subsystems=True): 

366 qmh = self.qm.spacing_def 

367 if corners is not None: 

368 v1 = np.array(corners[0]).ravel() / Bohr 

369 v2 = np.array(corners[1]).ravel() / Bohr 

370 else: # Use vacuum 

371 pos_old = atoms_in.get_positions()[0] 

372 dmy_atoms = atoms_in.copy() 

373 dmy_atoms.center(vacuum=vacuum) 

374 pos_new = dmy_atoms.get_positions()[0] 

375 v1 = (pos_old - pos_new) / Bohr 

376 v2 = v1 + np.diag(dmy_atoms.get_cell()) / Bohr 

377 

378 # Needed for restarting 

379 self.given_corner_v1 = v1 * Bohr 

380 self.given_corner_v2 = v2 * Bohr 

381 self.given_cell = atoms_in.get_cell() 

382 

383 # Sanity check: quantum box must be inside the classical one 

384 assert (all([v1[w] <= v2[w] and v1[w] >= 0 and 

385 v2[w] <= np.diag(self.cl.cell)[w] for w in range(3)])) 

386 

387 # Ratios of the user-given spacings 

388 self.hratios = self.cl.spacing_def / qmh 

389 self.num_refinements = \ 

390 1 + int(round(np.log(self.hratios[0]) / np.log(2.0))) 

391 assert ([int(round(np.log(self.hratios[w]) / np.log(2.0))) == 

392 self.num_refinements for w in range(3)]) 

393 

394 # Create quantum grid 

395 self.qm.cell = np.zeros((3, 3)) 

396 for w in range(3): 

397 self.qm.cell[w, w] = v2[w] - v1[w] 

398 

399 N_c = h2gpts(qmh, self.qm.cell, 4) 

400 # N_c = get_number_of_grid_points(self.qm.cell, qmh) 

401 self.qm.spacing = np.diag(self.qm.cell) / N_c 

402 

403 # Classical corner indices must be divisible with numb 

404 if any(self.cl.spacing / self.qm.spacing >= 3): 

405 numb = 1 

406 elif any(self.cl.spacing / self.qm.spacing >= 2): 

407 numb = 2 

408 else: 

409 numb = 4 

410 

411 # The index mismatch of the two simulation cells 

412 # Round before taking floor/ceil to avoid floating point 

413 # arithmetic errors 

414 num_indices_1 = numb * np.floor( 

415 np.round(np.array(v1) / self.cl.spacing / numb, 2)).astype(int) 

416 num_indices_2 = numb * np.ceil( 

417 np.round(np.array(v2) / self.cl.spacing / numb, 2)).astype(int) 

418 self.num_indices = num_indices_2 - num_indices_1 

419 

420 # Center, left, and right points of the suggested quantum grid 

421 cp = 0.5 * (np.array(v1) + np.array(v2)) 

422 lp = cp - 0.5 * self.num_indices * self.cl.spacing 

423 # rp = cp + 0.5 * self.num_indices * self.cl.spacing 

424 

425 # Indices in the classical grid restricting the quantum grid 

426 self.shift_indices_1 = np.round(lp / self.cl.spacing).astype(int) 

427 self.shift_indices_2 = self.shift_indices_1 + self.num_indices 

428 

429 # Sanity checks 

430 assert (all([self.shift_indices_1[w] >= 0 and 

431 self.shift_indices_2[w] <= self.cl.gd.N_c[w] 

432 for w in range(3)])), ( 

433 'Could not find appropriate quantum grid. ' 

434 'Move it further away from the boundary.') 

435 

436 # Corner coordinates 

437 self.qm.corner1 = self.shift_indices_1 * self.cl.spacing 

438 self.qm.corner2 = self.shift_indices_2 * self.cl.spacing 

439 

440 # Now the information for creating the subsystems is ready 

441 if create_subsystems: 

442 return self.create_subsystems(atoms_in) 

443 

444 def create_subsystems(self, atoms_in): 

445 

446 # Create new Atoms object 

447 atoms_out = atoms_in.copy() 

448 

449 # New quantum grid 

450 self.qm.cell = \ 

451 np.diag([(self.shift_indices_2[w] - self.shift_indices_1[w]) * 

452 self.cl.spacing[w] for w in range(3)]) 

453 self.qm.spacing = self.cl.spacing / self.hratios 

454 N_c = h2gpts(self.qm.spacing, self.qm.cell, 4) 

455 

456 atoms_out.set_cell(np.diag(self.qm.cell) * Bohr) 

457 atoms_out.positions = atoms_in.get_positions() - self.qm.corner1 * Bohr 

458 

459 msg = self.messages.append 

460 msg("Quantum box readjustment:") 

461 msg(" Given cell: [%10.5f %10.5f %10.5f]" % 

462 tuple(np.diag(atoms_in.get_cell()))) 

463 msg(" Given atomic coordinates:") 

464 for s, c in zip(atoms_in.get_chemical_symbols(), 

465 atoms_in.get_positions()): 

466 msg(" %s %10.5f %10.5f %10.5f" % 

467 (s, c[0], c[1], c[2])) 

468 msg(" Readjusted cell: [%10.5f %10.5f %10.5f]" % 

469 tuple(np.diag(atoms_out.get_cell()))) 

470 msg(" Readjusted atomic coordinates:") 

471 for s, c in zip(atoms_out.get_chemical_symbols(), 

472 atoms_out.get_positions()): 

473 msg(" %s %10.5f %10.5f %10.5f" % 

474 (s, c[0], c[1], c[2])) 

475 

476 msg(" Given corner points: " + 

477 "(%10.5f %10.5f %10.5f) - (%10.5f %10.5f %10.5f)" 

478 % (tuple(np.concatenate((self.given_corner_v1, 

479 self.given_corner_v2))))) 

480 msg(" Readjusted corner points: " + 

481 "(%10.5f %10.5f %10.5f) - (%10.5f %10.5f %10.5f)" 

482 % (tuple(np.concatenate((self.qm.corner1, 

483 self.qm.corner2)) * Bohr))) 

484 msg(" Indices in classical grid: " + 

485 "(%10i %10i %10i) - (%10i %10i %10i)" 

486 % (tuple(np.concatenate((self.shift_indices_1, 

487 self.shift_indices_2))))) 

488 msg(" Grid points in classical grid: " + 

489 "(%10i %10i %10i)" 

490 % (tuple(self.cl.gd.N_c))) 

491 msg(" Grid points in quantum grid: " + 

492 "(%10i %10i %10i)" 

493 % (tuple(N_c))) 

494 msg(" Spacings in quantum grid: " + 

495 "(%10.5f %10.5f %10.5f)" 

496 % (tuple(np.diag(self.qm.cell) * Bohr / N_c))) 

497 msg(" Spacings in classical grid: " + 

498 "(%10.5f %10.5f %10.5f)" 

499 % (tuple(np.diag(self.cl.cell) * 

500 Bohr / h2gpts(self.cl.spacing, self.cl.cell, 4)))) 

501 # msg(" Ratios of cl/qm spacings: " + 

502 # "(%10i %10i %10i)" % (tuple(self.hratios))) 

503 # msg(" = (%10.2f %10.2f %10.2f)" % 

504 # (tuple((np.diag(self.cl.cell) * Bohr / \ 

505 # get_number_of_grid_points(self.cl.cell, 

506 # self.cl.spacing)) / \ 

507 # (np.diag(self.qm.cell) * Bohr / N_c)))) 

508 msg(" Needed number of refinements: %10i" 

509 % self.num_refinements) 

510 

511 # First, create the quantum grid equivalent GridDescriptor 

512 # self.cl.subgd. Then coarsen it until its h_cv equals 

513 # that of self.cl.gd. Finally, map the points between 

514 # clgd and coarsened subgrid. 

515 subcell_cv = np.diag(self.qm.corner2 - self.qm.corner1) 

516 N_c = h2gpts(self.cl.spacing, subcell_cv, 4) 

517 N_c = self.shift_indices_2 - self.shift_indices_1 

518 self.cl.subgds = [] 

519 self.cl.subgds.append(GridDescriptor(N_c, subcell_cv, False, 

520 serial_comm, self.cl.dparsize)) 

521 

522 # msg(" N_c/spacing of the subgrid: " + 

523 # "%3i %3i %3i / %.4f %.4f %.4f" % 

524 # (self.cl.subgds[0].N_c[0], 

525 # self.cl.subgds[0].N_c[1], 

526 # self.cl.subgds[0].N_c[2], 

527 # self.cl.subgds[0].h_cv[0][0] * Bohr, 

528 # self.cl.subgds[0].h_cv[1][1] * Bohr, 

529 # self.cl.subgds[0].h_cv[2][2] * Bohr)) 

530 # msg(" shape from the subgrid: " + 

531 # "%3i %3i %3i" % (tuple(self.cl.subgds[0].empty().shape))) 

532 

533 self.cl.coarseners = [] 

534 self.cl.refiners = [] 

535 for n in range(self.num_refinements): 

536 self.cl.subgds.append(self.cl.subgds[n].refine()) 

537 self.cl.refiners.append(Transformer(self.cl.subgds[n], 

538 self.cl.subgds[n + 1])) 

539 

540 # msg(" refiners[%i] can perform the transformation " + 

541 # "(%3i %3i %3i) -> (%3i %3i %3i)" % (\ 

542 # n, 

543 # self.cl.subgds[n].empty().shape[0], 

544 # self.cl.subgds[n].empty().shape[1], 

545 # self.cl.subgds[n].empty().shape[2], 

546 # self.cl.subgds[n + 1].empty().shape[0], 

547 # self.cl.subgds[n + 1].empty().shape[1], 

548 # self.cl.subgds[n + 1].empty().shape[2])) 

549 self.cl.coarseners.append(Transformer(self.cl.subgds[n + 1], 

550 self.cl.subgds[n])) 

551 self.cl.coarseners[:] = self.cl.coarseners[::-1] 

552 

553 # Now extend the grid in order to handle 

554 # the zero boundary conditions that the refiner assumes 

555 # The default interpolation order 

556 self.extend_nn = \ 

557 Transformer(GridDescriptor([8, 8, 8], [1, 1, 1], 

558 False, serial_comm, None), 

559 GridDescriptor([8, 8, 8], [1, 1, 1], 

560 False, serial_comm, None).coarsen()).nn 

561 

562 self.extended_num_indices = self.num_indices + [2, 2, 2] 

563 

564 # Center, left, and right points of the suggested quantum grid 

565 extended_cp = 0.5 * (np.array(self.given_corner_v1 / Bohr) + 

566 np.array(self.given_corner_v2 / Bohr)) 

567 extended_lp = extended_cp - 0.5 * (self.extended_num_indices * 

568 self.cl.spacing) 

569 # extended_rp = extended_cp + 0.5 * (self.extended_num_indices * 

570 # self.cl.spacing) 

571 

572 # Indices in the classical grid restricting the quantum grid 

573 self.extended_shift_indices_1 = \ 

574 np.round(extended_lp / self.cl.spacing).astype(int) 

575 self.extended_shift_indices_2 = \ 

576 self.extended_shift_indices_1 + self.extended_num_indices 

577 

578 # msg(' extended_shift_indices_1: %i %i %i' 

579 # % (self.extended_shift_indices_1[0], 

580 # self.extended_shift_indices_1[1], 

581 # self.extended_shift_indices_1[2])) 

582 # msg(' extended_shift_indices_2: %i %i %i' 

583 # % (self.extended_shift_indices_2[0], 

584 # self.extended_shift_indices_2[1], 

585 # self.extended_shift_indices_2[2])) 

586 # msg(' cl.gd.N_c: %i %i %i' 

587 # % (self.cl.gd.N_c[0], self.cl.gd.N_c[1], self.cl.gd.N_c[2])) 

588 

589 # Sanity checks 

590 assert (all([self.extended_shift_indices_1[w] >= 0 and 

591 self.extended_shift_indices_2[w] <= self.cl.gd.N_c[w] 

592 for w in range(3)])), ( 

593 'Could not find appropriate quantum grid. ' 

594 'Move it further away from the boundary.') 

595 

596 # Corner coordinates 

597 self.qm.extended_corner1 = \ 

598 self.extended_shift_indices_1 * self.cl.spacing 

599 self.qm.extended_corner2 = \ 

600 self.extended_shift_indices_2 * self.cl.spacing 

601 N_c = self.extended_shift_indices_2 - self.extended_shift_indices_1 

602 

603 self.cl.extended_subgds = [] 

604 self.cl.extended_refiners = [] 

605 extended_subcell_cv = np.diag(self.qm.extended_corner2 - 

606 self.qm.extended_corner1) 

607 

608 self.cl.extended_subgds.append(GridDescriptor( 

609 N_c, extended_subcell_cv, False, serial_comm, None)) 

610 

611 for n in range(self.num_refinements): 

612 self.cl.extended_subgds.append(self.cl.extended_subgds[n].refine()) 

613 self.cl.extended_refiners.append(Transformer( 

614 self.cl.extended_subgds[n], self.cl.extended_subgds[n + 1])) 

615 

616 # msg(" extended_refiners[%i] can perform the transformation " + 

617 # "(%3i %3i %3i) -> (%3i %3i %3i)" 

618 # % (n, 

619 # self.cl.extended_subgds[n].empty().shape[0], 

620 # self.cl.extended_subgds[n].empty().shape[1], 

621 # self.cl.extended_subgds[n].empty().shape[2], 

622 # self.cl.extended_subgds[n + 1].empty().shape[0], 

623 # self.cl.extended_subgds[n + 1].empty().shape[1], 

624 # self.cl.extended_subgds[n + 1].empty().shape[2])) 

625 

626 # msg(" N_c/spacing of the refined subgrid: " + 

627 # "%3i %3i %3i / %.4f %.4f %.4f" 

628 # % (self.cl.subgds[-1].N_c[0], 

629 # self.cl.subgds[-1].N_c[1], 

630 # self.cl.subgds[-1].N_c[2], 

631 # self.cl.subgds[-1].h_cv[0][0] * Bohr, 

632 # self.cl.subgds[-1].h_cv[1][1] * Bohr, 

633 # self.cl.subgds[-1].h_cv[2][2] * Bohr)) 

634 # msg(" shape from the refined subgrid: %3i %3i %3i" 

635 # % (tuple(self.cl.subgds[-1].empty().shape))) 

636 

637 self.extended_deltaIndex = 2**(self.num_refinements) * self.extend_nn 

638 # msg(" self.extended_deltaIndex = %i" % self.extended_deltaIndex) 

639 

640 qgpts = self.cl.subgds[-1].coarsen().N_c 

641 

642 # Assure that one returns to the original shape 

643 dmygd = self.cl.subgds[-1].coarsen() 

644 for n in range(self.num_refinements - 1): 

645 dmygd = dmygd.coarsen() 

646 

647 # msg(" N_c/spacing of the coarsened subgrid: " + 

648 # "%3i %3i %3i / %.4f %.4f %.4f" 

649 # % (dmygd.N_c[0], dmygd.N_c[1], dmygd.N_c[2], 

650 # dmygd.h_cv[0][0] * Bohr, 

651 # dmygd.h_cv[1][1] * Bohr, 

652 # dmygd.h_cv[2][2] * Bohr)) 

653 

654 self.has_subsystems = True 

655 return atoms_out, self.qm.spacing[0] * Bohr, qgpts 

656 

657 def print_messages(self, printer_function): 

658 printer_function("\n *** QSFDTD ***\n") 

659 for msg in self.messages: 

660 printer_function(msg) 

661 

662 for msg in self.classical_material.messages: 

663 printer_function(msg) 

664 printer_function("\n *********************\n") 

665 

666 # Set the time step 

667 def set_time_step(self, time_step): 

668 self.time_step = time_step 

669 

670 # Set the time 

671 def set_time(self, time): 

672 self.time = time 

673 

674 # Setup kick 

675 def set_kick(self, kick): 

676 self.kick = np.array(kick) 

677 

678 def finalize_propagation(self): 

679 pass 

680 

681 def set_calculation_mode(self, calculation_mode): 

682 # Three calculation modes are available: 

683 # 1) solve: just solve the Poisson equation with 

684 # given quantum+classical rho 

685 # 2) iterate: iterate classical density so that the Poisson 

686 # equation for quantum+classical rho is satisfied 

687 # 3) propagate: propagate classical density in time, and solve 

688 # the new Poisson equation 

689 assert (calculation_mode == 'solve' or calculation_mode == 'iterate' or 

690 calculation_mode == 'propagate') 

691 self.calculation_mode = calculation_mode 

692 

693 # The density object must be attached, so that the electric field 

694 # from all-electron density can be calculated 

695 def set_density(self, density): 

696 self.density = density 

697 

698 # Returns the classical density and the grid descriptor 

699 def get_density(self, global_array=False): 

700 if global_array: 

701 return \ 

702 self.cl.gd.collect(self.classical_material.charge_density) * \ 

703 self.classical_material.sign, \ 

704 self.cl.gd 

705 else: 

706 return \ 

707 self.classical_material.charge_density * \ 

708 self.classical_material.sign, \ 

709 self.cl.gd 

710 

711 # Returns the quantum + classical density in the large classical box, 

712 # so that the classical charge is refined into it and the quantum 

713 # charge is coarsened there 

714 def get_combined_data(self, qmdata=None, cldata=None, spacing=None, 

715 qmgd=None, clgd=None): 

716 

717 if qmdata is None: 

718 qmdata = self.density.rhot_g 

719 

720 if cldata is None: 

721 cldata = (self.classical_material.charge_density * 

722 self.classical_material.sign) 

723 

724 if qmgd is None: 

725 qmgd = self.qm.gd 

726 

727 if clgd is None: 

728 clgd = self.cl.gd 

729 

730 if spacing is None: 

731 spacing = clgd.h_cv[0, 0] 

732 

733 spacing_au = spacing / Bohr # from Angstroms to a.u. 

734 

735 # Refine classical part 

736 clgd = clgd.new_descriptor() 

737 cl_g = cldata.copy() 

738 while clgd.h_cv[0, 0] > spacing_au * 1.50: # 45: 

739 clgd2 = clgd.refine() 

740 cl_g = Transformer(clgd, clgd2).apply(cl_g) 

741 clgd = clgd2 

742 

743 # Coarse quantum part 

744 qmgd = qmgd.new_descriptor() 

745 qm_g = qmdata.copy() 

746 while qmgd.h_cv[0, 0] < clgd.h_cv[0, 0] * 0.95: 

747 qmgd2 = qmgd.coarsen()