Source code for

from __future__ import annotations

from functools import cached_property
from typing import Any, Union

import numpy as np
from ase import Atoms
from ase.units import Bohr, Ha

from gpaw.core import UGDesc
from gpaw.core.atom_arrays import AtomDistribution
from gpaw.densities import Densities
from gpaw.electrostatic_potential import ElectrostaticPotential
from gpaw.gpu import as_np
from gpaw.mpi import broadcast_float, world
from import trace, zips
from import Density
from import IBZWaveFunctions
from import InputParameters
from import Logger
from import Potential
from import SCFLoop
from gpaw.setup import Setups
from gpaw.typing import Array1D, Array2D
from gpaw.utilities import (check_atoms_too_close,
from gpaw.utilities.partition import AtomPartition

class ReuseWaveFunctionsError(Exception):
    """Reusing the old wave functions after cell-change failed.

    Most likekly, the number of k-points changed.

class NonsenseError(Exception):
    """Operation doesn't make sense."""

class CalculationModeError(Exception):
    """Calculation mode does not match what is expected from a given method.

    For example, if a method only works in collinear mode and receives a
    calculator in non-collinear mode, this exception should be raised.

units = {'energy': Ha,
         'free_energy': Ha,
         'forces': Ha / Bohr,
         'stress': Ha / Bohr**3,
         'dipole': Bohr,
         'magmom': 1.0,
         'magmoms': 1.0,
         'non_collinear_magmom': 1.0,
         'non_collinear_magmoms': 1.0}

[docs]class DFTState: def __init__(self, ibzwfs: IBZWaveFunctions, density: Density, potential: Potential): """State of a Kohn-Sham calculation.""" self.ibzwfs = ibzwfs self.density = density self.potential = potential def __repr__(self): return (f'DFTState({self.ibzwfs!r}, ' f'{self.density!r}, {self.potential!r})') def __str__(self): return f'{self.ibzwfs}\n{self.density}\n{self.potential}'
[docs] def move(self, fracpos_ac, atomdist): self.ibzwfs.move(fracpos_ac, atomdist) self.density.move(fracpos_ac, atomdist) self.potential.move(atomdist)
[docs]class DFTCalculation: def __init__(self, state: DFTState, setups: Setups, scf_loop: SCFLoop, pot_calc, log: Logger): self.state = state self.setups = setups self.scf_loop = scf_loop self.pot_calc = pot_calc self.log = log self.comm = log.comm self.results: dict[str, Any] = {} self.fracpos_ac = self.pot_calc.fracpos_ac
[docs] @classmethod def from_parameters(cls, atoms: Atoms, params: Union[dict, InputParameters], comm=None, log=None, builder=None) -> DFTCalculation: """Create DFTCalculation object from parameters and atoms.""" from import builder as create_builder check_atoms_too_close(atoms) check_atoms_too_close_to_boundary(atoms) if params is None: params = {} if isinstance(params, dict): params = InputParameters(params) if not isinstance(log, Logger): log = Logger(log, comm or world) builder = builder or create_builder(atoms, params, log.comm) basis_set = builder.create_basis_set() density = builder.density_from_superposition(basis_set) density.normalize() # The SCF-loop has a Hamiltonian that has an fft-plan that is # cached for later use, so best to create the SCF-loop first # FIX this! scf_loop = builder.create_scf_loop() pot_calc = builder.create_potential_calculator() potential, _ = pot_calc.calculate_without_orbitals( density, kpt_band_comm=builder.communicators['D']) ibzwfs = builder.create_ibz_wave_functions( basis_set, potential, log=log) if ibzwfs.wfs_qs[0][0]._eig_n is not None: ibzwfs.calculate_occs(scf_loop.occ_calc) state = DFTState(ibzwfs, density, potential) write_atoms(atoms, builder.initial_magmom_av, builder.grid, log) log(state) log(builder.setups) log(scf_loop) log(pot_calc) return cls(state, builder.setups, scf_loop, pot_calc, log)
[docs] def move_atoms(self, atoms) -> DFTCalculation: check_atoms_too_close(atoms) self.fracpos_ac = np.ascontiguousarray(atoms.get_scaled_positions()) self.comm.broadcast(self.fracpos_ac, 0) atomdist = self.state.density.D_asii.layout.atomdist grid = self.state.density.nt_sR.desc rank_a = grid.ranks_from_fractional_positions(self.fracpos_ac) atomdist = AtomDistribution(rank_a, atomdist.comm) self.pot_calc.move(self.fracpos_ac, atomdist) self.state.move(self.fracpos_ac, atomdist) mm_av = self.results['non_collinear_magmoms'] write_atoms(atoms, mm_av, self.state.density.nt_sR.desc, self.log) self.results = {} return self
[docs] def iconverge(self, convergence=None, maxiter=None, calculate_forces=None): self.state.ibzwfs.make_sure_wfs_are_read_from_gpw_file() yield from self.scf_loop.iterate(self.state, self.pot_calc, convergence, maxiter, calculate_forces, log=self.log)
[docs] @trace def converge(self, convergence=None, maxiter=None, steps=99999999999999999, calculate_forces=None): """Converge to self-consistent solution of Kohn-Sham equation.""" for step, _ in enumerate(self.iconverge(convergence, maxiter, calculate_forces), start=1): if step == steps: break else: # no break self.log(scf_steps=step)
[docs] def energies(self): energies = combine_energies(self.state.potential, self.state.ibzwfs) self.log('Energy contributions relative to reference atoms:', f'(reference = {self.setups.Eref * Ha:.6f})\n') for name, e in energies.items(): if not name.startswith('total') and name != 'stress': self.log(f'{name + ":":10} {e * Ha:14.6f}') total_free = energies['total_free'] total_extrapolated = energies['total_extrapolated'] self.log('----------------------------') self.log(f'Free energy: {total_free * Ha:14.6f}') self.log(f'Extrapolated:{total_extrapolated * Ha:14.6f}\n') total_free = broadcast_float(total_free, self.comm) total_extrapolated = broadcast_float(total_extrapolated, self.comm) self.results['free_energy'] = total_free self.results['energy'] = total_extrapolated
[docs] def dipole(self): if 'dipole' in self.results: return dipole_v = self.state.density.calculate_dipole_moment(self.fracpos_ac) x, y, z = dipole_v * Bohr self.log(f'dipole moment: [{x:.6f}, {y:.6f}, {z:.6f}] # |e|*Ang\n') self.results['dipole'] = dipole_v
[docs] def magmoms(self) -> tuple[Array1D, Array2D]: mm_v, mm_av = self.state.density.calculate_magnetic_moments() self.results['magmom'] = mm_v[2] self.results['magmoms'] = mm_av[:, 2].copy() self.results['non_collinear_magmoms'] = mm_av self.results['non_collinear_magmom'] = mm_v if self.state.density.ncomponents > 1: x, y, z = mm_v self.log(f'total magnetic moment: [{x:.6f}, {y:.6f}, {z:.6f}]\n') self.log('local magnetic moments: [') for a, (setup, m_v) in enumerate(zips(self.setups, mm_av)): x, y, z = m_v c = ',' if a < len(mm_av) - 1 else ']' self.log(f' [{x:9.6f}, {y:9.6f}, {z:9.6f}]{c}' f' # {setup.symbol:2} {a}') self.log() return mm_v, mm_av
[docs] def forces(self, silent=False): """Calculate atomic forces.""" if 'forces' not in self.results or silent: self._calculate_forces() if silent: return self.log('\nforces: [ # eV/Ang') F_av = self.results['forces'] * (Ha / Bohr) for a, setup in enumerate(self.setups): x, y, z = F_av[a] c = ',' if a < len(F_av) - 1 else ']' self.log(f' [{x:9.3f}, {y:9.3f}, {z:9.3f}]{c}' f' # {setup.symbol:2} {a}')
def _calculate_forces(self): xc = self.pot_calc.xc assert not hasattr(xc.xc, 'setup_force_corrections') # Force from projector functions (and basis set): F_av = self.state.ibzwfs.forces(self.state.potential) pot_calc = self.pot_calc Fcc_avL, Fnct_av, Ftauct_av, Fvbar_av = pot_calc.force_contributions( self.state) # Force from compensation charges: ccc_aL = \ self.state.density.calculate_compensation_charge_coefficients() for a, dF_vL in Fcc_avL.items(): F_av[a] += dF_vL @ ccc_aL[a] # Force from smooth core charge: for a, dF_v in Fnct_av.items(): F_av[a] += dF_v[:, 0] if Ftauct_av is not None: # Force from smooth core ked: for a, dF_v in Ftauct_av.items(): F_av[a] += dF_v[:, 0] # Force from zero potential: for a, dF_v in Fvbar_av.items(): F_av[a] += dF_v[:, 0] F_av = as_np(F_av) domain_comm = ccc_aL.layout.atomdist.comm domain_comm.sum(F_av) F_av = self.state.ibzwfs.ibz.symmetries.symmetrize_forces(F_av) self.comm.broadcast(F_av, 0) self.results['forces'] = F_av
[docs] def stress(self) -> None: if 'stress' in self.results: return stress_vv = self.pot_calc.stress(self.state) self.log('\nstress tensor: [ # eV/Ang^3') for (x, y, z), c in zips(stress_vv * (Ha / Bohr**3), ',,]'): self.log(f' [{x:13.6f}, {y:13.6f}, {z:13.6f}]{c}') self.results['stress'] = stress_vv.flat[[0, 4, 8, 5, 2, 1]]
[docs] def write_converged(self) -> None: self.state.ibzwfs.write_summary(self.log) vacuum_level = self.state.potential.get_vacuum_level() if not np.isnan(vacuum_level): self.log(f'vacuum-level: {vacuum_level:.3f} # V') try: wf1, wf2 = self.workfunctions(vacuum_level=vacuum_level) except NonsenseError: pass else: self.log(f'Workfunctions: {wf1:.3f}, {wf2:.3f} # eV') self.log.fd.flush()
[docs] def workfunctions(self, *, vacuum_level: float | None = None ) -> tuple[float, float]: if vacuum_level is None: vacuum_level = self.state.potential.get_vacuum_level() if np.isnan(vacuum_level): raise NonsenseError('No vacuum') try: correction = self.pot_calc.poisson_solver.dipole_layer_correction() except NotImplementedError: raise NonsenseError('No dipole layer') correction *= Ha fermi_level = self.state.ibzwfs.fermi_level * Ha wf = vacuum_level - fermi_level return wf - correction, wf + correction
[docs] def electrostatic_potential(self) -> ElectrostaticPotential: return ElectrostaticPotential.from_calculation(self)
[docs] def densities(self) -> Densities: return Densities.from_calculation(self)
@cached_property def _atom_partition(self): # Backwards compatibility helper atomdist = self.state.density.D_asii.layout.atomdist return AtomPartition(atomdist.comm, atomdist.rank_a)
[docs] def new(self, atoms: Atoms, params: InputParameters, log=None) -> DFTCalculation: """Create new DFTCalculation object.""" from import builder as create_builder if params.mode['name'] != 'pw': raise ReuseWaveFunctionsError ibzwfs = self.state.ibzwfs if ibzwfs.domain_comm.size != 1: raise ReuseWaveFunctionsError if not self.state.density.nt_sR.desc.pbc_c.all(): raise ReuseWaveFunctionsError check_atoms_too_close(atoms) check_atoms_too_close_to_boundary(atoms) builder = create_builder(atoms, params, self.comm) kpt_kc = builder.ibz.kpt_kc old_kpt_kc = ibzwfs.ibz.kpt_kc if len(kpt_kc) != len(old_kpt_kc): raise ReuseWaveFunctionsError if abs(kpt_kc - old_kpt_kc).max() > 1e-9: raise ReuseWaveFunctionsError log('# Interpolating wave functions to new cell') density =, builder.interpolation_desc, builder.fracpos_ac, builder.atomdist) density.normalize() # Make sure all have exactly the same density. # Not quite sure it is needed??? # At the moment we skip it on GPU's because it doesn't # work! if density.nt_sR.xp is np: self.comm.broadcast(, 0) scf_loop = builder.create_scf_loop() pot_calc = builder.create_potential_calculator() potential, _ = pot_calc.calculate(density) old_ibzwfs = ibzwfs def create_wfs(spin, q, k, kpt_c, weight): wfs = old_ibzwfs.wfs_qs[q][spin] return wfs.morph( builder.wf_desc, builder.fracpos_ac, builder.atomdist) ibzwfs = ibzwfs.create( ibz=builder.ibz, nelectrons=old_ibzwfs.nelectrons, ncomponents=old_ibzwfs.ncomponents, create_wfs_func=create_wfs, kpt_comm=old_ibzwfs.kpt_comm, kpt_band_comm=old_ibzwfs.kpt_band_comm, comm=self.comm) state = DFTState(ibzwfs, density, potential) write_atoms(atoms, builder.initial_magmom_av, builder.grid, log) log(state) log(builder.setups) log(scf_loop) log(pot_calc) return DFTCalculation( state, builder.setups, scf_loop, pot_calc, log)
def combine_energies(potential: Potential, ibzwfs: IBZWaveFunctions) -> dict[str, float]: """Add up energy contributions.""" energies = potential.energies.copy() energies.pop('stress', 0.0) energies['kinetic'] += ibzwfs.energies['band'] energies['kinetic'] += ibzwfs.energies.get('exx_kinetic', 0.0) energies['xc'] += (ibzwfs.energies.get('exx_vv', 0.0) + ibzwfs.energies.get('exx_vc', 0.0) + ibzwfs.energies.get('exx_cc', 0.0)) energies['entropy'] = ibzwfs.energies['entropy'] energies['total_free'] = sum(energies.values()) energies['total_extrapolated'] = (energies['total_free'] + ibzwfs.energies['extrapolation']) return energies def write_atoms(atoms: Atoms, magmom_av: Array2D, grid: UGDesc, log) -> None: from gpaw.output import print_cell, print_positions print_positions(atoms, log, magmom_av) print_cell(grid._gd, grid.pbc, log)