from __future__ import annotations
import warnings
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 UGArray, 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 as bcast
from gpaw.mpi import broadcast_float, world
from gpaw.new import trace, zips
from gpaw.new.density import Density
from gpaw.new.energies import DFTEnergies
from gpaw.new.ibzwfs import IBZWaveFunctions
from gpaw.new.input_parameters import InputParameters
from gpaw.new.logger import Logger
from gpaw.new.potential import Potential
from gpaw.new.scf import SCFLoop
from gpaw.setup import Setups
from gpaw.typing import Array1D, Array2D
from gpaw.utilities import (check_atoms_too_close,
check_atoms_too_close_to_boundary)
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 DFTCalculation:
def __init__(self,
atoms: Atoms,
ibzwfs: IBZWaveFunctions,
density: Density,
potential: Potential,
setups: Setups,
scf_loop: SCFLoop,
pot_calc,
log: Logger,
params: InputParameters,
energies: DFTEnergies | None = None):
self.atoms = atoms
self.ibzwfs = ibzwfs
self.density = density
self.potential = potential
self.setups = setups
self.scf_loop = scf_loop
self.pot_calc = pot_calc
self.log = log
self.comm = log.comm
self.params = params
self.results: dict[str, Any] = {}
self.relpos_ac = self.pot_calc.relpos_ac
self.energies = energies or DFTEnergies()
self.forces_have_been_printed = False
[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 gpaw.new.builder 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, log)
basis_set = builder.create_basis_set()
# 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(log)
density = builder.density_from_superposition(basis_set)
if len(atoms) == 0:
density.nt_sR.data[:] = 1.0
density.normalize(pot_calc.environment.charge)
potential, energies, _ = 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:
nelectrons = (density.nvalence - density.charge +
pot_calc.environment.charge)
ibzwfs.calculate_occs(scf_loop.occ_calc, nelectrons)
write_atoms(atoms, builder.initial_magmom_av, builder.grid, log)
log(ibzwfs)
log(density)
log(potential)
log(builder.setups)
log(scf_loop)
log(pot_calc)
return cls(atoms, ibzwfs, density, potential,
builder.setups, scf_loop, pot_calc, log,
params=params, energies=energies)
[docs]
def get_ase_calc(self):
"""Create ASE-compatible GPAW calculator.
"""
from gpaw.new.ase_interface import ASECalculator
return ASECalculator(params=self.params,
log=self.log,
dft=self,
atoms=self.atoms)
[docs]
def move_atoms(self, atoms) -> DFTCalculation:
check_atoms_too_close(atoms)
self.atoms = atoms
self.relpos_ac = np.ascontiguousarray(atoms.get_scaled_positions())
self.comm.broadcast(self.relpos_ac, 0)
atomdist = self.density.D_asii.layout.atomdist
grid = self.density.nt_sR.desc
rank_a = grid.ranks_from_fractional_positions(self.relpos_ac)
atomdist = AtomDistribution(rank_a, atomdist.comm)
self.pot_calc.move(self.relpos_ac, atomdist)
self.ibzwfs.move(self.relpos_ac, atomdist)
self.density.move(self.relpos_ac, atomdist)
if self.ibzwfs.has_wave_functions():
self.density.update(self.ibzwfs)
self.potential.move(atomdist)
self.potential, self.energies, _ = self.pot_calc.calculate(
self.density, self.ibzwfs, self.potential.vHt_x)
mm_av = self.results['non_collinear_magmoms']
write_atoms(atoms, mm_av, self.density.nt_sR.desc, self.log)
self.results = {}
self.forces_have_been_printed = False
return self
[docs]
def iconverge(self, maxiter=None, calculate_forces=None):
self.ibzwfs.make_sure_wfs_are_read_from_gpw_file()
for ctx in self.scf_loop.iterate(self.ibzwfs,
self.density,
self.potential,
self.energies,
self.pot_calc,
maxiter=maxiter,
calculate_forces=calculate_forces,
log=self.log):
self.energies = ctx.energies
self.potential = ctx.potential
yield ctx
[docs]
@trace
def converge(self,
maxiter=None,
steps=99999999999999999,
calculate_forces=None):
"""Converge to self-consistent solution of Kohn-Sham equation."""
for step, _ in enumerate(self.iconverge(maxiter,
calculate_forces),
start=1):
if step == steps:
break
else: # no break
self.log('SCF steps:', step)
[docs]
def energy(self):
self.results['free_energy'] = broadcast_float(
self.energies.total_free, self.comm)
self.results['energy'] = broadcast_float(
self.energies.total_extrapolated, self.comm)
self.log('Energy contributions relative to reference atoms:',
f'(reference = {self.setups.Eref * Ha:.6f})\n')
self.energies.summary(self.log)
[docs]
def dipole(self):
if 'dipole' in self.results:
return
dipole_v = self.density.calculate_dipole_moment(self.relpos_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.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.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):
"""Calculate atomic forces."""
if 'forces' not in self.results:
self._calculate_forces()
if self.forces_have_been_printed:
return
self.forces_have_been_printed = True
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:10.4f}, {y:10.4f}, {z:10.4f}]{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.ibzwfs.forces(self.potential)
pot_calc = self.pot_calc
Q_aL = self.density.calculate_compensation_charge_coefficients()
Fcc_av, Fnct_av, Ftauct_av, Fvbar_av, Fext_av = \
pot_calc.force_contributions(Q_aL,
self.density,
self.potential)
# Force from compensation charges:
F_av += Fcc_av
# 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 = Q_aL.layout.atomdist.comm
domain_comm.sum(F_av)
# Force from extensions (only from rank 0)
F_av += Fext_av
F_av = self.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.ibzwfs, self.density, self.potential)
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.ibzwfs.write_summary(self.log)
vacuum_level = self.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 (top, bottom)')
self.log.fd.flush()
[docs]
def workfunctions(self,
*, _vacuum_level=None) -> tuple[float, float]:
if _vacuum_level is None:
_vacuum_level = self.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.ibzwfs.fermi_level * Ha
wf = _vacuum_level - fermi_level
return wf - correction, wf + correction
[docs]
def vacuum_level(self) -> float:
return self.potential.get_vacuum_level()
[docs]
def electrostatic_potential(self) -> ElectrostaticPotential:
return ElectrostaticPotential.from_calculation(self)
[docs]
def densities(self) -> Densities:
return Densities.from_calculation(self)
[docs]
def wave_function(self, band: int, kpt=0, spin=None,
periodic=False,
broadcast=True) -> UGArray:
psit_nR = self.wave_functions(n1=band, n2=band + 1, kpt=kpt, spin=spin,
periodic=periodic, broadcast=broadcast)
if psit_nR is not None:
return psit_nR[0]
[docs]
def wave_functions(self, n1=0, n2=None, kpt=0, spin=None,
periodic=False,
broadcast=True,
_pad=True) -> UGArray:
collinear = self.ibzwfs.collinear
if collinear:
if spin is None:
spin = 0
else:
assert spin is None or spin == 0
wfs = self.ibzwfs.get_wfs(spin=spin if collinear else 0,
kpt=kpt,
n1=n1, n2=n2)
if wfs is not None:
basis = getattr(self.scf_loop.hamiltonian, 'basis', None)
grid = self.density.nt_sR.desc.new(comm=None)
if collinear:
wfs = wfs.to_uniform_grid_wave_functions(grid, basis)
psit_nR = wfs.psit_nX
else:
psit_nsG = wfs.psit_nX
grid = grid.new(kpt=psit_nsG.desc.kpt_c,
dtype=psit_nsG.desc.dtype)
psit_nR = psit_nsG.ifft(grid=grid)
if not psit_nR.desc.pbc.all() and _pad:
psit_nR = psit_nR.to_pbc_grid()
if periodic:
psit_nR.multiply_by_eikr(-psit_nR.desc.kpt_c)
else:
psit_nR = None
if broadcast:
psit_nR = bcast(psit_nR, 0, self.comm)
return psit_nR.scaled(cell=Bohr, values=Bohr**-1.5)
@cached_property
def _atom_partition(self):
# Backwards compatibility helper
atomdist = self.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 gpaw.new.builder import builder as create_builder
if params.mode['name'] != 'pw':
raise ReuseWaveFunctionsError
ibzwfs = self.ibzwfs
if ibzwfs.domain_comm.size != 1:
raise ReuseWaveFunctionsError
if not self.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, log)
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')
scf_loop = builder.create_scf_loop()
pot_calc = builder.create_potential_calculator(log)
density = self.density.new(builder.grid,
builder.interpolation_desc,
builder.relpos_ac,
builder.atomdist)
density.normalize(pot_calc.environment.charge)
# 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(density.nt_sR.data, 0)
potential, energies, _ = 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.relpos_ac,
builder.atomdist)
ibzwfs = ibzwfs.create(
ibz=builder.ibz,
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)
write_atoms(atoms, builder.initial_magmom_av, builder.grid, log)
log(ibzwfs)
log(density)
log(potential)
log(builder.setups)
log(scf_loop)
log(pot_calc)
return DFTCalculation(
atoms, ibzwfs, density, potential,
builder.setups, scf_loop, pot_calc, log,
params=params, energies=energies)
[docs]
def get_state(self):
return DFTState(self.ibzwfs, self.density, self.potential)
@property
def state(self):
warnings.warn('Use of deprecated DFTCalculation.state attribute. '
'Use ibzwfs, density and potential attributes instead.')
return self.get_state()
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)
[docs]
class DFTState:
def __init__(self,
ibzwfs: IBZWaveFunctions,
density: Density,
potential: Potential,
energies):
"""State of a Kohn-Sham calculation."""
self.ibzwfs = ibzwfs
self.density = density
self.potential = potential
self.energies = energies