Coverage for gpaw/new/calculation.py: 88%
297 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1from __future__ import annotations
3import warnings
4from functools import cached_property
5from typing import Any, TYPE_CHECKING
7import numpy as np
8from ase import Atoms
9from ase.units import Bohr, Ha
10from gpaw.core import UGArray, UGDesc
11from gpaw.core.atom_arrays import AtomDistribution
12from gpaw.densities import Densities
13from gpaw.electrostatic_potential import ElectrostaticPotential
14from gpaw.gpu import as_np
15from gpaw.mpi import broadcast as bcast
16from gpaw.mpi import broadcast_float, world
17from gpaw.new import trace, zips
18from gpaw.new.density import Density
19from gpaw.new.energies import DFTEnergies
20from gpaw.new.ibzwfs import IBZWaveFunctions
21from gpaw.new.logger import Logger
22from gpaw.new.potential import Potential
23from gpaw.new.scf import SCFLoop
24from gpaw.setup import Setups
25from gpaw.typing import Array1D, Array2D
26from gpaw.utilities import (check_atoms_too_close,
27 check_atoms_too_close_to_boundary)
28from gpaw.utilities.partition import AtomPartition
29if TYPE_CHECKING:
30 from gpaw.dft import Parameters
33class ReuseWaveFunctionsError(Exception):
34 """Reusing the old wave functions after cell-change failed.
36 Most likekly, the number of k-points changed.
37 """
40class NonsenseError(Exception):
41 """Operation doesn't make sense."""
44class CalculationModeError(Exception):
45 """Calculation mode does not match what is expected from a given method.
47 For example, if a method only works in collinear mode and receives a
48 calculator in non-collinear mode, this exception should be raised.
49 """
52units = {'energy': Ha,
53 'free_energy': Ha,
54 'forces': Ha / Bohr,
55 'stress': Ha / Bohr**3,
56 'dipole': Bohr,
57 'magmom': 1.0,
58 'magmoms': 1.0,
59 'non_collinear_magmom': 1.0,
60 'non_collinear_magmoms': 1.0}
63class DFTCalculation:
64 def __init__(self,
65 atoms: Atoms,
66 ibzwfs: IBZWaveFunctions,
67 density: Density,
68 potential: Potential,
69 setups: Setups,
70 scf_loop: SCFLoop,
71 pot_calc,
72 log: Logger,
73 params: Parameters,
74 energies: DFTEnergies | None = None):
75 self.atoms = atoms
76 self.ibzwfs = ibzwfs
77 self.density = density
78 self.potential = potential
79 self.setups = setups
80 self.scf_loop = scf_loop
81 self.pot_calc = pot_calc
82 self.log = log
83 self.comm = log.comm
84 self.params = params
86 self.results: dict[str, Any] = {}
87 self.relpos_ac = self.pot_calc.relpos_ac
88 self.energies = energies or DFTEnergies()
89 self.forces_have_been_printed = False
91 def __getattr__(self, name):
92 matches = [ext
93 for ext in self.pot_calc.extensions
94 if ext.name == name]
95 if len(matches) != 1:
96 raise AttributeError
97 return matches[0]
99 @classmethod
100 def from_parameters(cls,
101 atoms: Atoms,
102 params: Parameters,
103 comm=None,
104 log=None) -> DFTCalculation:
105 """Create DFTCalculation object from parameters and atoms."""
106 check_atoms_too_close(atoms)
107 check_atoms_too_close_to_boundary(atoms)
109 if not isinstance(log, Logger):
110 log = Logger(log, comm or world)
112 builder = params.dft_component_builder(atoms, log=log)
114 basis_set = builder.create_basis_set()
116 # The SCF-loop has a Hamiltonian that has an fft-plan that is
117 # cached for later use, so best to create the SCF-loop first
118 # FIX this!
119 scf_loop = builder.create_scf_loop()
121 pot_calc = builder.create_potential_calculator()
123 density = builder.density_from_superposition(basis_set)
124 if len(atoms) == 0:
125 density.nt_sR.data[:] = 1.0
126 density.normalize(pot_calc.environment.charge)
128 potential, energies, _ = pot_calc.calculate_without_orbitals(
129 density, kpt_band_comm=builder.communicators['D'])
130 ibzwfs = builder.create_ibz_wave_functions(
131 basis_set, potential)
133 if ibzwfs.wfs_qs[0][0]._eig_n is not None:
134 nelectrons = (density.nvalence - density.charge +
135 pot_calc.environment.charge)
136 ibzwfs.calculate_occs(scf_loop.occ_calc, nelectrons)
138 write_atoms(atoms, builder.initial_magmom_av, builder.grid, log)
139 log(ibzwfs)
140 log(density)
141 log(potential)
142 log(builder.setups)
143 log(scf_loop)
144 log(pot_calc)
146 return cls(atoms, ibzwfs, density, potential,
147 builder.setups, scf_loop, pot_calc, log,
148 params=params, energies=energies)
150 def ase_calculator(self):
151 """Create ASE-compatible GPAW calculator.
152 """
153 from gpaw.new.ase_interface import ASECalculator
154 return ASECalculator(params=self.params,
155 log=self.log,
156 dft=self,
157 atoms=self.atoms)
159 def move_atoms(self, atoms) -> DFTCalculation:
160 check_atoms_too_close(atoms)
162 self.atoms = atoms
163 self.relpos_ac = np.ascontiguousarray(atoms.get_scaled_positions())
164 self.comm.broadcast(self.relpos_ac, 0)
166 atomdist = self.density.D_asii.layout.atomdist
167 grid = self.density.nt_sR.desc
168 rank_a = grid.ranks_from_fractional_positions(self.relpos_ac)
169 atomdist = AtomDistribution(rank_a, atomdist.comm)
171 self.pot_calc.move(self.relpos_ac, atomdist)
172 self.ibzwfs.move(self.relpos_ac, atomdist)
173 self.density.move(self.relpos_ac, atomdist)
174 if self.ibzwfs.has_wave_functions():
175 self.density.update(self.ibzwfs)
176 self.potential.move(atomdist)
178 self.potential, self.energies, _ = self.pot_calc.calculate(
179 self.density, self.ibzwfs, self.potential.vHt_x)
181 mm_av = self.results['non_collinear_magmoms']
182 write_atoms(atoms, mm_av, self.density.nt_sR.desc, self.log)
184 self.results = {}
185 self.forces_have_been_printed = False
187 return self
189 def iconverge(self, maxiter=None, calculate_forces=None):
190 self.ibzwfs.make_sure_wfs_are_read_from_gpw_file()
191 for ctx in self.scf_loop.iterate(self.ibzwfs,
192 self.density,
193 self.potential,
194 self.energies,
195 self.pot_calc,
196 maxiter=maxiter,
197 calculate_forces=calculate_forces,
198 log=self.log):
199 self.energies = ctx.energies
200 self.potential = ctx.potential
201 yield ctx
203 @trace
204 def converge(self,
205 maxiter=None,
206 steps=99999999999999999,
207 calculate_forces=None):
208 """Converge to self-consistent solution of Kohn-Sham equation."""
209 for step, _ in enumerate(self.iconverge(maxiter,
210 calculate_forces),
211 start=1):
212 if step == steps:
213 break
214 else: # no break
215 self.log('SCF steps:', step)
217 def energy(self):
218 self.results['free_energy'] = broadcast_float(
219 self.energies.total_free, self.comm)
220 self.results['energy'] = broadcast_float(
221 self.energies.total_extrapolated, self.comm)
223 self.log('Energy contributions relative to reference atoms:',
224 f'(reference = {self.setups.Eref * Ha:.6f})\n')
225 self.energies.summary(self.log)
227 def dipole(self):
228 if 'dipole' in self.results:
229 return
230 dipole_v = self.density.calculate_dipole_moment(self.relpos_ac)
231 x, y, z = dipole_v * Bohr
232 self.log(f'dipole moment: [{x:.6f}, {y:.6f}, {z:.6f}] # |e|*Ang\n')
233 self.results['dipole'] = dipole_v
235 def magmoms(self) -> tuple[Array1D, Array2D]:
236 mm_v, mm_av = self.density.calculate_magnetic_moments()
237 self.results['magmom'] = mm_v[2]
238 self.results['magmoms'] = mm_av[:, 2].copy()
239 self.results['non_collinear_magmoms'] = mm_av
240 self.results['non_collinear_magmom'] = mm_v
242 if self.density.ncomponents > 1:
243 x, y, z = mm_v
244 self.log(f'total magnetic moment: [{x:.6f}, {y:.6f}, {z:.6f}]\n')
245 self.log('local magnetic moments: [')
246 for a, (setup, m_v) in enumerate(zips(self.setups, mm_av)):
247 x, y, z = m_v
248 c = ',' if a < len(mm_av) - 1 else ']'
249 self.log(f' [{x:9.6f}, {y:9.6f}, {z:9.6f}]{c}'
250 f' # {setup.symbol:2} {a}')
251 self.log()
252 return mm_v, mm_av
254 def forces(self):
255 """Calculate atomic forces."""
256 if 'forces' not in self.results:
257 self._calculate_forces()
259 if self.forces_have_been_printed:
260 return
262 self.forces_have_been_printed = True
263 self.log('\nForces: [ # eV/Ang')
264 F_av = self.results['forces'] * (Ha / Bohr)
265 for a, setup in enumerate(self.setups):
266 x, y, z = F_av[a]
267 c = ',' if a < len(F_av) - 1 else ']'
268 self.log(f' [{x:10.4f}, {y:10.4f}, {z:10.4f}]{c}'
269 f' # {setup.symbol:2} {a}')
271 def _calculate_forces(self):
272 xc = self.pot_calc.xc
273 assert not hasattr(xc.xc, 'setup_force_corrections')
275 # Force from projector functions (and basis set):
276 F_av = self.ibzwfs.forces(self.potential)
278 pot_calc = self.pot_calc
279 Q_aL = self.density.calculate_compensation_charge_coefficients()
280 Fcc_av, Fnct_av, Ftauct_av, Fvbar_av, Fext_av = \
281 pot_calc.force_contributions(Q_aL,
282 self.density,
283 self.potential)
285 # Force from compensation charges:
286 F_av += Fcc_av
288 # Force from smooth core charge:
289 for a, dF_v in Fnct_av.items():
290 F_av[a] += dF_v[:, 0]
292 if Ftauct_av is not None:
293 # Force from smooth core ked:
294 for a, dF_v in Ftauct_av.items():
295 F_av[a] += dF_v[:, 0]
297 # Force from zero potential:
298 for a, dF_v in Fvbar_av.items():
299 F_av[a] += dF_v[:, 0]
301 F_av = as_np(F_av)
303 domain_comm = Q_aL.layout.atomdist.comm
304 domain_comm.sum(F_av)
306 # Force from extensions (only from rank 0)
307 F_av += Fext_av
309 F_av = self.ibzwfs.ibz.symmetries.symmetrize_forces(F_av)
310 self.comm.broadcast(F_av, 0)
311 self.results['forces'] = F_av
313 def stress(self) -> None:
314 if 'stress' in self.results:
315 return
316 stress_vv = self.pot_calc.stress(
317 self.ibzwfs, self.density, self.potential)
318 self.log('\nstress tensor: [ # eV/Ang^3')
319 for (x, y, z), c in zips(stress_vv * (Ha / Bohr**3), ',,]'):
320 self.log(f' [{x:13.6f}, {y:13.6f}, {z:13.6f}]{c}')
321 self.results['stress'] = stress_vv.flat[[0, 4, 8, 5, 2, 1]]
323 def write_converged(self) -> None:
324 self.ibzwfs.write_summary(self.log)
325 vacuum_level = self.potential.get_vacuum_level()
326 if not np.isnan(vacuum_level):
327 self.log(f'vacuum-level: {vacuum_level:.3f} # V')
328 try:
329 wf1, wf2 = self.workfunctions(_vacuum_level=vacuum_level)
330 except NonsenseError:
331 pass
332 else:
333 self.log(
334 f'Workfunctions: {wf1:.3f}, {wf2:.3f} eV (top, bottom)')
335 self.log.fd.flush()
337 def workfunctions(self,
338 *, _vacuum_level=None) -> tuple[float, float]:
339 if _vacuum_level is None:
340 _vacuum_level = self.potential.get_vacuum_level()
341 if np.isnan(_vacuum_level):
342 raise NonsenseError('No vacuum')
343 try:
344 correction = self.pot_calc.poisson_solver.dipole_layer_correction()
345 except NotImplementedError:
346 raise NonsenseError('No dipole layer')
347 correction *= Ha
348 fermi_level = self.ibzwfs.fermi_level * Ha
349 wf = _vacuum_level - fermi_level
350 return wf - correction, wf + correction
352 def vacuum_level(self) -> float:
353 return self.potential.get_vacuum_level()
355 def electrostatic_potential(self) -> ElectrostaticPotential:
356 return ElectrostaticPotential.from_calculation(self)
358 def densities(self) -> Densities:
359 return Densities.from_calculation(self)
361 def wave_function(self, band: int, kpt=0, spin=None,
362 periodic=False,
363 broadcast=True) -> UGArray:
364 psit_nR = self.wave_functions(n1=band, n2=band + 1, kpt=kpt, spin=spin,
365 periodic=periodic, broadcast=broadcast)
366 if psit_nR is not None:
367 return psit_nR[0]
369 def wave_functions(self, n1=0, n2=None, kpt=0, spin=None,
370 periodic=False,
371 broadcast=True,
372 _pad=True) -> UGArray:
373 collinear = self.ibzwfs.collinear
374 if collinear:
375 if spin is None:
376 spin = 0
377 else:
378 assert spin is None or spin == 0
379 wfs = self.ibzwfs.get_wfs(spin=spin if collinear else 0,
380 kpt=kpt,
381 n1=n1, n2=n2)
382 if wfs is not None:
383 basis = getattr(self.scf_loop.hamiltonian, 'basis', None)
384 grid = self.density.nt_sR.desc.new(comm=None)
385 if collinear:
386 wfs = wfs.to_uniform_grid_wave_functions(grid, basis)
387 psit_nR = wfs.psit_nX
388 else:
389 psit_nsG = wfs.psit_nX
390 grid = grid.new(kpt=psit_nsG.desc.kpt_c,
391 dtype=psit_nsG.desc.dtype)
392 psit_nR = psit_nsG.ifft(grid=grid)
393 if not psit_nR.desc.pbc.all() and _pad:
394 psit_nR = psit_nR.to_pbc_grid()
395 if periodic:
396 psit_nR.multiply_by_eikr(-psit_nR.desc.kpt_c)
397 else:
398 psit_nR = None
399 if broadcast:
400 psit_nR = bcast(psit_nR, 0, self.comm)
401 return psit_nR.scaled(cell=Bohr, values=Bohr**-1.5)
403 @cached_property
404 def _atom_partition(self):
405 # Backwards compatibility helper
406 atomdist = self.density.D_asii.layout.atomdist
407 return AtomPartition(atomdist.comm, atomdist.rank_a)
409 def new(self,
410 atoms: Atoms,
411 params: Parameters,
412 log=None) -> DFTCalculation:
413 """Create new DFTCalculation object."""
415 if params.mode.name != 'pw':
416 raise ReuseWaveFunctionsError
418 ibzwfs = self.ibzwfs
419 # if ibzwfs.domain_comm.size != 1:
420 # raise ReuseWaveFunctionsError
422 check_atoms_too_close(atoms)
423 check_atoms_too_close_to_boundary(atoms)
425 builder = params.dft_component_builder(atoms, log=log)
427 kpt_kc = builder.ibz.kpt_kc
428 old_kpt_kc = ibzwfs.ibz.kpt_kc
429 if len(kpt_kc) != len(old_kpt_kc):
430 raise ReuseWaveFunctionsError
431 if abs(kpt_kc - old_kpt_kc).max() > 1e-9:
432 raise ReuseWaveFunctionsError
434 log('Interpolating wave functions to new cell')
436 scf_loop = builder.create_scf_loop()
437 pot_calc = builder.create_potential_calculator()
439 density = self.density.new(builder.grid,
440 builder.interpolation_desc,
441 builder.relpos_ac,
442 builder.atomdist)
443 density.normalize(pot_calc.environment.charge)
445 # Make sure all have exactly the same density.
446 # Not quite sure it is needed???
447 # At the moment we skip it on GPU's because it doesn't
448 # work!
449 if density.nt_sR.xp is np:
450 ibzwfs.kpt_band_comm.broadcast(density.nt_sR.data, 0)
452 potential, energies, _ = pot_calc.calculate(density)
454 old_ibzwfs = ibzwfs
456 def create_wfs(spin, q, k, kpt_c, weight):
457 wfs = old_ibzwfs.wfs_qs[q][spin]
458 return wfs.morph(
459 builder.wf_desc,
460 builder.relpos_ac,
461 builder.atomdist)
463 ibzwfs = ibzwfs.create(
464 ibz=builder.ibz,
465 ncomponents=old_ibzwfs.ncomponents,
466 create_wfs_func=create_wfs,
467 kpt_comm=old_ibzwfs.kpt_comm,
468 kpt_band_comm=old_ibzwfs.kpt_band_comm,
469 comm=self.comm)
471 write_atoms(atoms, builder.initial_magmom_av, builder.grid, log)
472 log(ibzwfs)
473 log(density)
474 log(potential)
475 log(builder.setups)
476 log(scf_loop)
477 log(pot_calc)
479 return DFTCalculation(
480 atoms, ibzwfs, density, potential,
481 builder.setups, scf_loop, pot_calc, log,
482 params=params, energies=energies)
484 def get_state(self):
485 return DFTState(self.ibzwfs, self.density, self.potential)
487 @property
488 def state(self):
489 warnings.warn('Use of deprecated DFTCalculation.state attribute. '
490 'Use ibzwfs, density and potential attributes instead.')
491 return self.get_state()
494def write_atoms(atoms: Atoms,
495 magmom_av: Array2D,
496 grid: UGDesc,
497 log) -> None:
498 from gpaw.output import print_cell, print_positions
499 print_positions(atoms, log, magmom_av)
500 print_cell(grid._gd, grid.pbc, log)
503class DFTState:
504 def __init__(self,
505 ibzwfs: IBZWaveFunctions,
506 density: Density,
507 potential: Potential,
508 energies):
509 """State of a Kohn-Sham calculation."""
510 self.ibzwfs = ibzwfs
511 self.density = density
512 self.potential = potential
513 self.energies = energies