"""Module to handle ab initio electrostatic potentials from the DFT code GPAW."""
from __future__ import annotations
import io
import threading
import warnings
from collections import defaultdict
from dataclasses import dataclass
from functools import partial
from numbers import Number
from typing import Any, List, Tuple, Union
import dask
import dask.array as da
import numpy as np
from ase import Atoms, units
from ase.data import atomic_numbers, chemical_symbols
from dask.delayed import Delayed
from scipy.interpolate import interp1d
from abtem.atoms import is_cell_orthogonal, plane_to_axes
from abtem.core.axes import AxisMetadata
from abtem.core.backend import copy_to_device
from abtem.core.electron_configurations import (
config_str_to_config_tuples,
electron_configurations,
)
from abtem.core.ensemble import _wrap_with_array
from abtem.core.fft import fft_crop
from abtem.core.utils import itemset
from abtem.inelastic.phonons import (
BaseFrozenPhonons,
DummyFrozenPhonons,
FrozenPhonons,
_safe_read_atoms,
)
from abtem.parametrizations import EwaldParametrization, LobatoParametrization
from abtem.potentials.charge_density import _interpolate_slice
from abtem.potentials.iam import Potential, PotentialArray, _PotentialBuilder
from abtem.slicing import _integrate_planes
try:
from gpaw import GPAW
except ImportError:
GPAW = None
if GPAW is not None:
from gpaw.atom.shapefunc import shape_functions
from gpaw.mpi import SerialCommunicator
# GPAW keeps one FFT plan per grid shape and dtype for the whole process
# (``gpaw.fftw.create_plans`` caches them in a module-level dict), and each plan
# owns the scratch arrays (``tmp_R``, ``tmp_Q``) that every transform writes its
# input to and reads its output from. All calculators on the same grid share
# these arrays, so two threads evaluating GPAW at once overwrite each other's
# transforms: a read gives a wrong potential without an error, and concurrent
# ``_DummyGPAW.setups`` calls can make GPAW itself fail. Dask's threaded
# scheduler runs the tasks of a multi-file ``GPAWPotential`` in parallel, and
# tasks of an earlier build can still be running when the caller's thread
# returns to GPAW. Every abTEM call that makes a GPAW calculator compute on its
# grids (reading a ``.gpw`` file, initializing, evaluating a density or
# potential) therefore holds this lock; attribute reads, and GPAW's radial atom
# solvers, which create no FFT plans, do not. The lock is re-entrant because
# ``_read_gpw`` calls ``_DummyGPAW.from_gpaw``.
_GPAW_LOCK = threading.RLock()
@dataclass
class _DummyGPAW:
setup_mode: str
setup_xc: str
nt_sG: np.ndarray
gd: Any
D_asp: np.ndarray
atoms: Atoms
Q_aL: dict
valence_potential: np.ndarray
def __len__(self):
return 1
@property
def setups(self):
with _GPAW_LOCK:
gpaw = GPAW(txt=None, mode=self.setup_mode, xc=self.setup_xc)
gpaw.initialize(self.atoms)
return gpaw.setups
@classmethod
def from_gpaw(cls, gpaw, lazy: bool = True):
# if lazy:
# return dask.delayed(cls.from_gpaw)(gpaw, lazy=False)
with _GPAW_LOCK:
return cls._from_gpaw(gpaw)
@classmethod
def _from_gpaw(cls, gpaw):
atoms = gpaw.atoms.copy()
atoms.calc = None
valence_potential = gpaw.get_electrostatic_potential()
# GPAW's new-style calculator (default since GPAW 26) renames the
# compensation charge coefficients from Q_aL to ccc_aL. Read it after
# get_electrostatic_potential(), which fills in Q_aL on a restarted
# old-style calculator.
Q_aL = getattr(gpaw.density, "Q_aL", None)
if Q_aL is None:
Q_aL = gpaw.density.ccc_aL
kwargs = {
"setup_mode": gpaw.parameters.mode,
"setup_xc": gpaw.parameters.xc,
"nt_sG": gpaw.density.nt_sG.copy(),
"gd": gpaw.density.gd.new_descriptor(comm=SerialCommunicator()),
"D_asp": dict(gpaw.density.D_asp),
"atoms": atoms,
"valence_potential": valence_potential,
"Q_aL": dict(Q_aL),
}
return cls(**kwargs)
@classmethod
def from_file(cls, path: str, lazy: bool = True):
if lazy:
return dask.delayed(_read_gpw)(path)
return _read_gpw(path)
@classmethod
def from_generic(cls, calculator, lazy: bool = True):
if isinstance(calculator, str):
return cls.from_file(calculator, lazy=lazy)
elif hasattr(calculator, "density"):
return cls.from_gpaw(calculator, lazy=lazy)
elif isinstance(calculator, cls):
return calculator
else:
raise RuntimeError()
def _read_gpw(path: str) -> _DummyGPAW:
# Module level, so that a task calling it pickles it by reference; a
# classmethod in a task is pickled by value together with the globals it
# uses, and the lock cannot be pickled.
with _GPAW_LOCK:
return _DummyGPAW.from_gpaw(GPAW(path))
class _LastRead:
# Reads the ``.gpw`` placeholders of one eager ``_partition_args`` call and
# keeps only the calculator read last. Consecutive blocks that share a
# placeholder (one path with ``FrozenPhonons``) reuse it; a block with
# another placeholder drops it before reading. An eager build therefore
# holds one read calculator while it generates slices, and two while it
# reads the next file, since the caller still holds the previous block's
# potential. The object lives in that call's block array and goes with
# it, so no calculator outlives the build.
def __init__(self):
self._lock = threading.Lock()
self._placeholder = None
self._calculator = None
def read(self, placeholder: Delayed) -> "_DummyGPAW":
with self._lock:
if placeholder is not self._placeholder:
self._placeholder = None
self._calculator = None
# Synchronous: one task, and no thread pool started from
# inside a caller's task. ``_read_gpw`` takes the GPAW lock.
self._calculator = placeholder.compute(scheduler="synchronous")
self._placeholder = placeholder
return self._calculator
class _PendingRead:
# A ``.gpw`` placeholder in an eager block, read when the block's potential
# is made (``GPAWPotential._gpaw_potential``), not when the blocks are.
__slots__ = ("placeholder", "reads")
def __init__(self, placeholder: Delayed, reads: _LastRead):
self.placeholder = placeholder
self.reads = reads
def _defer_placeholders(calculators, reads: _LastRead):
# ``from_file`` leaves each ``.gpw`` path as a ``Delayed`` read. A lazy
# graph resolves it as a task argument; eager blocks read it one block at
# a time through ``reads``.
if isinstance(calculators, list):
return [_defer_placeholders(calculator, reads) for calculator in calculators]
if isinstance(calculators, Delayed):
return _PendingRead(calculators, reads)
return calculators
def _read_pending(calculators):
if isinstance(calculators, list):
return [_read_pending(calculator) for calculator in calculators]
if isinstance(calculators, _PendingRead):
return calculators.reads.read(calculators.placeholder)
return calculators
[docs]
def get_core_correction_interpolators(setups, D_asp, Q_aL, rcgauss):
interpolators = []
for a, D_sp in D_asp.items():
setup = setups[a]
c = setup.xc_correction
rgd = c.rgd
params = setup.data.shape_function.copy()
params["lmax"] = 0
ghat_g = shape_functions(rgd, **params)[0]
Z_g = shape_functions(rgd, "gauss", rcgauss, lmax=0)[0] * setup.Z
D_q = np.dot(D_sp.sum(0), c.B_pqL[:, :, 0])
dn_g = np.dot(D_q, (c.n_qg - c.nt_qg)) * np.sqrt(4 * np.pi)
dn_g += 4 * np.pi * (c.nc_g - c.nct_g)
dn_g -= Z_g
dn_g -= Q_aL[a][0] * ghat_g * np.sqrt(4 * np.pi)
dv_g = rgd.poisson(dn_g) / np.sqrt(4 * np.pi)
dv_g[1:] /= rgd.r_g[1:]
dv_g[0] = dv_g[1]
dv_g[-1] = 0.0
interpolator = interp1d(
units.Bohr * rgd.r_g,
-dv_g / np.sqrt(4 * np.pi) * units.Ha,
fill_value=0,
bounds_error=False,
)
interpolators.append(interpolator)
return interpolators
def _same_elements_and_cell(atoms, other):
"""Whether two atoms have the same elements in the same order and cell."""
return (
len(atoms) == len(other)
and np.array_equal(atoms.numbers, other.numbers)
and np.allclose(np.array(atoms.cell), np.array(other.cell))
)
[docs]
def integrate_slice(array, gpts, a, b, thickness):
"""
Integrate the planes of `array` along its last axis between the heights `a`
and `b`, and Fourier interpolate the result to `gpts`.
Plane k is the sample at the height k * dz, with dz the plane spacing, and
stands for the heights from k * dz - dz / 2 to k * dz + dz / 2; the grid is
periodic, so plane 0 also stands for the heights within dz / 2 below the
thickness. A slice takes each plane times the length of the plane's heights
that lie between `a` and `b`. Consecutive slices share their limit, so the
weights of every plane add up to dz. A limit on a plane splits it between the
slices on either side (the trapezoidal rule), so the integral of a slice is
centred on the slice. Taking the planes from floor(a / dz) to floor(b / dz)
whole (a left Riemann sum) put it dz / 2 below the slice.
"""
nz = array.shape[2]
dz = thickness / nz
# The limits in plane spacings; the grid ends at plane nz.
start, stop = a / dz, min(b / dz, nz)
if stop <= start:
return np.zeros(gpts, dtype=array.dtype)
slice_array = _integrate_planes(array, start, stop) * dz
old_shape = slice_array.shape
slice_array = np.fft.fftn(slice_array)
slice_array = fft_crop(slice_array, gpts)
slice_array = np.fft.ifftn(slice_array).real * np.prod(gpts) / np.prod(old_shape)
return slice_array
class _DummyParametrization:
def __init__(self, potential):
self._potential = potential
super().__init__()
def potential(self, symbol):
return self._potential
@property
def sigmas(self):
return {}
def _generate_slices(
interpolators,
valence_potential,
atoms,
gpts,
slice_thickness,
plane="xy",
first_slice=0,
last_slice=None,
device=None,
):
# `atoms` is the calculator's atoms or whole copies of them one after another
# (`atoms * repetitions`), so atom i of every copy, atoms[i::n], takes the core
# correction of atom i of the calculator.
n = len(interpolators)
potentials = []
for i, interpolator in enumerate(interpolators):
parametrization = _DummyParametrization(interpolator)
potential = Potential(
gpts=gpts,
atoms=atoms[i::n],
parametrization=parametrization,
slice_thickness=slice_thickness,
projection="finite",
plane=plane,
device=device,
)
potentials.append(potential)
if last_slice is None:
last_slice = len(potential)
potential_generators = [
potential.generate_slices(first_slice, last_slice) for potential in potentials
]
transform_valence_potential = None
if potential.plane != "xy":
if not is_cell_orthogonal(atoms.cell):
raise NotImplementedError
axes = plane_to_axes(potential.plane)
valence_potential = np.moveaxis(valence_potential, axes[:2], (0, 1))
transform_valence_potential = False
# else:
# atoms = ewald_potential.frozen_phonons.atoms
transformed_atoms = potential.get_transformed_atoms()
if np.allclose(transformed_atoms.cell, atoms.cell):
transform_valence_potential = False
elif transform_valence_potential is None:
transform_valence_potential = True
for i, slice_idx in enumerate(range(first_slice, last_slice)):
slic = next(potential_generators[0])
a, b = potential.get_sliced_atoms().slice_limits[slice_idx]
for potential_generator in potential_generators[1:]:
slic.array[:] += next(potential_generator).array
# The valence potential is on the host; its slice is moved to the device
# of the slice array.
if transform_valence_potential:
valence_slice = _interpolate_slice(
valence_potential, atoms.cell, potential.gpts, potential.sampling, a, b
)
else:
valence_slice = integrate_slice(
valence_potential, potential.gpts, a, b, potential.thickness
)
slic.array[:] -= copy_to_device(valence_slice, slic.array)
yield slic
[docs]
class GPAWPotential(_PotentialBuilder):
"""
Calculate the electrostatic potential from a (set of) converged GPAW DFT
calculation(s). Frozen phonons can be included either by specifying multiple GPAW
calculators corresponding to the different phonon configurations, or approximately
for a single calculator by using the `frozen_phonons` keyword.
Parameters
----------
calculators : (list of) gpaw.calculator.GPAW or (list of) str
GPAW calculator or path to GPAW calculator or list of calculators or paths.
Atoms are read from the calculator.
gpts : one or two int, optional
Number of grid points in `x` and `y` describing each slice of the potential.
Provide either "sampling" (spacing between consecutive grid points) or "gpts"
(total number of grid points).
sampling : one or two float, optional
Sampling of the potential in `x` and `y` [1 / Å].
Provide either "sampling" or "gpts".
slice_thickness : float or sequence of float, optional
Thickness of the potential slices in the propagation direction in [Å]
(default is 0.5 Å).
If given as a float the number of slices is calculated by dividing the slice
thickness into the `z`-height of supercell. The slice thickness may be given as
a sequence of values for each slice, in which case an error will be thrown if
the sum of slice thicknesses is not equal to the height of the atoms.
exit_planes : int or tuple of int, optional
The `exit_planes` argument can be used to calculate thickness series.
Providing `exit_planes` as a tuple of int indicates that the tuple contains the
slice indices after which an exit plane is desired, and hence during a
multislice simulation a measurement is created. If `exit_planes` is an integer a
measurement will be collected every `exit_planes` number of slices.
plane : str or two tuples of three float, optional
The plane relative to the provided atoms mapped to `xy` plane of the potential,
i.e. provided plane is perpendicular to the propagation direction. If string,
it must be a concatenation of two of 'x', 'y' and 'z';
the default value 'xy' indicates that potential slices are cuts along the
`xy`-plane of the atoms. The plane may also be specified with two arbitrary
3D vectors, which are mapped to the `x` and `y` directions of the potential.
The length of the vectors has no influence. If the vectors are not
perpendicular, the second vector is rotated in the plane to become perpendicular
to the first. Providing a value of ((1., 0., 0.), (0., 1., 0.)) is equivalent to
providing 'xy'.
origin : three float, optional
The origin relative to the provided Atoms mapped to the origin of the Potential.
This is equivalent to translating the atoms. The default is (0., 0., 0.)
box : three float, optional
The extent of the potential in `x`, `y` and `z`. If not given this is determined
from the atoms' cell. If the box size does not match an integer number of the
atoms' supercell, an affine transformation may be necessary to preserve
periodicity, determined by the `periodic` keyword
periodic : bool
If a transformation of the atomic structure is required, `periodic` determines
how the atomic structure is transformed. If True (default), the periodicity of
the Atoms is preserved, which may require applying a small affine
transformation to the atoms. If False, the transformed potential is effectively
cut out of a larger repeated potential, which may not preserve periodicity.
frozen_phonons : abtem.AbstractFrozenPhonons, optional
Approximates frozen phonons for a single GPAW calculator by displacing only the
nuclear core potentials. Supercedes the atoms from the calculator. A
trajectory (:class:`~abtem.inelastic.phonons.AtomsEnsemble`) is not accepted:
pass one calculator per frame instead.
repetitions : tuple of int
Repeats the atoms by integer amounts in the `x`, `y` and `z` directions before
applying frozen phonon displacements to calculate the potential contribution of
the nuclear cores. Necessary when using frozen phonons.
gridrefinement : int
Necessary interpolation of the charge density into a finer grid for improved
numerical precision. Allowed values are '2' and '4'.
device : str, optional
The device used for calculating the potential, 'cpu' or 'gpu'. The default is
determined by the user configuration file.
"""
[docs]
def __init__(
self,
calculators: Union["GPAW", List["GPAW"], List[str], str],
gpts: Union[int, Tuple[int, int]] = None,
sampling: Union[float, Tuple[float, float]] = None,
slice_thickness: float = 1.0,
exit_planes: int = None,
plane: str = "xy",
origin: Tuple[float, float, float] = (0.0, 0.0, 0.0),
box: Tuple[float, float, float] = None,
periodic: bool = True,
frozen_phonons: BaseFrozenPhonons = None,
repetitions: Tuple[int, int, int] = (1, 1, 1),
gridrefinement: int = 4,
device: str = None,
):
if GPAW is None:
raise RuntimeError(
"This functionality of abTEM requires GPAW, see https://wiki.fysik.dtu.dk/gpaw/."
)
if isinstance(calculators, (tuple, list)):
atoms = _safe_read_atoms(calculators[0])
num_configs = len(calculators)
if frozen_phonons is not None:
raise ValueError()
calculators = [
_DummyGPAW.from_generic(calculator) for calculator in calculators
]
frozen_phonons = DummyFrozenPhonons(atoms, num_configs=num_configs)
else:
atoms = _safe_read_atoms(calculators)
calculators = _DummyGPAW.from_generic(calculators)
if frozen_phonons is None:
frozen_phonons = DummyFrozenPhonons(atoms, num_configs=None)
elif frozen_phonons.ensemble_shape and not isinstance(
frozen_phonons, FrozenPhonons
):
raise ValueError(
"One calculator takes `FrozenPhonons`; for a trajectory, pass "
"one calculator per frame."
)
elif not _same_elements_and_cell(frozen_phonons.atoms, atoms):
raise ValueError(
"The frozen phonons must have the calculator's atoms: the same "
"elements in the same order and the same cell, with positions "
"that may differ. Repeat the cell with `repetitions`."
)
self._calculators = calculators
self._frozen_phonons = frozen_phonons
self._gridrefinement = gridrefinement
self._repetitions = tuple(repetitions)
cell = frozen_phonons.atoms.cell * repetitions
frozen_phonons.atoms.calc = None
super().__init__(
array_object=PotentialArray,
gpts=gpts,
sampling=sampling,
cell=cell,
slice_thickness=slice_thickness,
exit_planes=exit_planes,
device=device,
plane=plane,
origin=origin,
box=box,
periodic=periodic,
)
@property
def frozen_phonons(self):
return self._frozen_phonons
@property
def num_configurations(self) -> int:
"""Number of potential configurations; 1 when there is no ensemble axis."""
return len(self.frozen_phonons)
@property
def repetitions(self):
return self._repetitions
@property
def gridrefinement(self):
return self._gridrefinement
@property
def calculators(self):
return self._calculators
def _get_all_electron_density(self):
calculator = (
self.calculators[0]
if isinstance(self.calculators, list)
else self.calculators
)
try:
calculator = calculator.compute()
except AttributeError:
pass
# assert len(self.calculators) == 1
calculator = _DummyGPAW.from_generic(calculator)
atoms = self.frozen_phonons.atoms
if self.repetitions != (1, 1, 1):
cell_cv = calculator.gd.cell_cv * self.repetitions
N_c = tuple(
n_c * rep for n_c, rep in zip(calculator.gd.N_c, self.repetitions)
)
gd = calculator.gd.new_descriptor(N_c=N_c, cell_cv=cell_cv)
atoms = atoms * self.repetitions
nt_sG = np.tile(calculator.nt_sG, self.repetitions)
else:
gd = calculator.gd
nt_sG = calculator.nt_sG
random_atoms = self.frozen_phonons.randomize(atoms)
calc = GPAW(txt=None, mode=calculator.setup_mode, xc=calculator.setup_xc)
calc.initialize(random_atoms)
return self._get_all_electron_density(
nt_sG=nt_sG,
gd=gd,
D_asp=calculator.D_asp,
setups=calc.setups,
gridrefinement=self.gridrefinement,
atoms=random_atoms,
)
def _get_ewald_potential(self):
ewald_parametrization = EwaldParametrization(width=3)
atoms = self.frozen_phonons.atoms * self.repetitions
atoms = self.frozen_phonons.randomize(atoms)
ewald_potential = Potential(
atoms=atoms,
gpts=self.gpts,
sampling=self.sampling,
parametrization=ewald_parametrization,
slice_thickness=self.slice_thickness,
projection="finite",
# integral_method="quadrature",
plane=self.plane,
box=self.box,
origin=self.origin,
exit_planes=self.exit_planes,
device=self.device,
)
return ewald_potential
[docs]
def generate_slices(self, first_slice: int = 0, last_slice: int = None):
"""
Generate the slices for the potential.
Parameters
----------
first_slice : int, optional
Index of the first slice of the generated potential.
last_slice : int, optional
Index of the last slice of the generated potential.
Returns
-------
slices : generator of np.ndarray
Generator for the array of slices.
"""
if last_slice is None:
last_slice = len(self)
calculator = (
self.calculators[0]
if isinstance(self.calculators, list)
else self.calculators
)
try:
calculator = calculator.compute()
except AttributeError:
pass
# assert len(self.calculators) == 1
calculator = _DummyGPAW.from_generic(calculator)
atoms = self.frozen_phonons.atoms
valence_potential = calculator.valence_potential
if self.repetitions != (1, 1, 1):
# The valence potential is periodic with the calculator's cell.
valence_potential = np.tile(valence_potential, self.repetitions)
# cell_cv = calculator.gd.cell_cv * self.repetitions
# N_c = tuple(
# n_c * rep for n_c, rep in zip(calculator.gd.N_c, self.repetitions)
# )
# gd = calculator.gd.new_descriptor(N_c=N_c, cell_cv=cell_cv)
atoms = atoms * self.repetitions
# nt_sG = np.tile(calculator.nt_sG, self.repetitions)
# else:
# gd = calculator.gd
# nt_sG = calculator.nt_sG
random_atoms = self.frozen_phonons.randomize(atoms)
interpolators = get_core_correction_interpolators(
calculator.setups, calculator.D_asp, calculator.Q_aL, 0.001
)
# calc = GPAW(txt=None, mode=calculator.setup_mode, xc=calculator.setup_xc)
# calc.initialize(random_atoms)
# ewald_potential = self._get_ewald_potential()
# array = self._get_all_electron_density()
# array = calculator.valence_potential
for slic in _generate_slices(
interpolators,
plane=self.plane,
valence_potential=valence_potential,
atoms=random_atoms,
gpts=self.gpts,
slice_thickness=self.slice_thickness,
first_slice=first_slice,
last_slice=last_slice,
device=self.device,
):
yield slic
# for slic in _generate_slices(
# array, ewald_potential, first_slice=first_slice, last_slice=last_slice
# ):
# yield slic
@property
def ensemble_axes_metadata(self) -> List[AxisMetadata]:
return self._frozen_phonons.ensemble_axes_metadata
@property
def num_frozen_phonons(self):
return self.num_configurations
@property
def ensemble_shape(self):
return self._frozen_phonons.ensemble_shape
@staticmethod
def _gpaw_potential(*args, frozen_phonons_partial, **kwargs):
args = args[0]
if hasattr(args, "item"):
args = args.item()
if args["frozen_phonons"] is not None:
frozen_phonons = frozen_phonons_partial(args["frozen_phonons"])
frozen_phonons = frozen_phonons.item()
else:
frozen_phonons = None
calculators = _read_pending(args["calculators"])
new_potential = GPAWPotential(
calculators, frozen_phonons=frozen_phonons, **kwargs
)
return _wrap_with_array(new_potential)
def _from_partitioned_args(self):
kwargs = self._copy_kwargs(exclude=("calculators", "frozen_phonons"))
frozen_phonons_partial = self.frozen_phonons._from_partitioned_args()
return partial(
self._gpaw_potential,
frozen_phonons_partial=frozen_phonons_partial,
**kwargs,
)
def _partition_args(self, chunks: int = 1, lazy: bool = True):
chunks = self._validate_ensemble_chunks(chunks)
def frozen_phonons(calculators, frozen_phonons):
arr = np.zeros((1,), dtype=object)
itemset(
arr, 0, {"calculators": calculators, "frozen_phonons": frozen_phonons}
)
return arr
calculators = self.calculators
if not lazy:
calculators = _defer_placeholders(calculators, _LastRead())
if isinstance(self.frozen_phonons, FrozenPhonons):
array = np.zeros(len(self.frozen_phonons), dtype=object)
partitioned = self.frozen_phonons._partition_args(chunks, lazy=lazy)[0]
# Iterating a dask array yields per-element dask array slices;
# passing those into dask.delayed() and wrapping the result with
# another from_delayed()/concatenate() confuses dask's shape
# inference (IndexError deep in dask.array.core during compute).
# Delayed objects from to_delayed() compose correctly instead.
fp_chunks = partitioned.to_delayed().ravel() if lazy else partitioned
for i, fp in enumerate(fp_chunks):
if lazy:
block = dask.delayed(frozen_phonons)(calculators, fp)
itemset(array, i, da.from_delayed(block, shape=(1,), dtype=object))
else:
itemset(array, i, frozen_phonons(calculators, fp))
if lazy:
array = da.concatenate(list(array))
else:
if len(self.ensemble_shape) == 0 and lazy:
block = dask.delayed(frozen_phonons)(calculators, None)
array = da.from_delayed(block, shape=(), dtype=object)
elif len(self.ensemble_shape) == 0 and not lazy:
array = frozen_phonons(calculators, None)
array = _wrap_with_array(array, ndims=0)
else:
array = np.zeros(self.ensemble_shape[0], dtype=object)
for i, calculator in enumerate(calculators):
calculator = [calculator]
if lazy:
calculator = dask.delayed(calculator)
block = da.from_delayed(
dask.delayed(frozen_phonons)(calculator, None),
shape=(1,),
dtype=object,
)
else:
block = frozen_phonons(calculator, None)
itemset(array, i, block)
if lazy:
array = da.concatenate(list(array))
return (array,)
# return (array,)
# # if lazy:
# # print(array)
# # return (da.concatenate(list(array)),)
# # else:
# # return (array,)
[docs]
class GPAWParametrization:
"""
Calculate an Independent Atomic Model (IAM) potential based on a GPAW DFT
calculation.
"""
[docs]
def __init__(self, nodes=None, integration_step=0.002):
self._nodes = nodes
self._integration_step = integration_step
self._potential_functions = {}
def _get_added_electrons(self, symbol, charge):
if not charge:
return []
charge = int(np.sign(charge) * np.ceil(np.abs(charge)))
number = atomic_numbers[symbol]
config = config_str_to_config_tuples(
electron_configurations[chemical_symbols[number]]
)
ionic_config = config_str_to_config_tuples(
electron_configurations[chemical_symbols[number - charge]]
)
config = defaultdict(lambda: 0, {shell[:2]: shell[2] for shell in config})
ionic_config = defaultdict(
lambda: 0, {shell[:2]: shell[2] for shell in ionic_config}
)
electrons = []
for key in set(config.keys()).union(set(ionic_config.keys())):
difference = config[key] - ionic_config[key]
for i in range(np.abs(difference)):
electrons.append(key + (-np.sign(difference),))
return electrons
def _get_all_electron_atom(self, symbol, charge=0.0):
if isinstance(symbol, Number):
symbol = chemical_symbols[symbol]
added_electrons = self._get_added_electrons(symbol, charge)
try:
return self._run_all_electron_atom(symbol, added_electrons, spinpol=True)
except AssertionError:
# GPAW's radial solver fails to find a bound state for some open-shell
# atoms (e.g. Pm, Sm, Eu, Pu) in the spin-polarized calculation,
# independent of the SCF mixing. The spin-polarized and unpolarized
# total densities differ by only ~0.1 electrons, which is negligible
# for the independent-atom potential.
warnings.warn(
f"The spin-polarized all-electron calculation for {symbol} failed "
"to converge in GPAW; falling back to a non-spin-polarized "
"calculation."
)
return self._run_all_electron_atom(symbol, added_electrons, spinpol=False)
@staticmethod
def _run_all_electron_atom(symbol, added_electrons, spinpol):
from gpaw.atom.aeatom import AllElectronAtom
ae = AllElectronAtom(symbol, spinpol=spinpol, xc="PBE", log=io.StringIO())
for n, l, df in added_electrons:
ae.add(n, l, df)
if added_electrons:
# The occupations perturbed by add() need a more conservative
# mixing schedule to reach self-consistency than the default,
# which is tuned for the unperturbed (neutral) configuration.
ae.run(mix=0.005, maxiter=5000, dnmax=1e-5)
else:
ae.run()
ae.scalar_relativistic = True
ae.refine()
return ae
[docs]
def charge(self, symbol: str, charge: float = 0.0):
"""
Calculate the radial charge density for an atom.
Parameters
----------
symbol : str
Chemical symbol of the atomic element.
charge : float, optional
Charge the atom by the given fractional number of electrons.
Returns
-------
charge : callable
Function of the radial charge density with parameter 'r' corresponding to
the radial distance from the core.
"""
ae = self._get_all_electron_atom(symbol, charge)
r = ae.rgd.r_g * units.Bohr
n = ae.n_sg.sum(0) / units.Bohr**3
return interp1d(r, n, fill_value="extrapolate", bounds_error=False)
def x_ray_scattering_factor(self, symbol: str, charge: float = 0.0):
if isinstance(symbol, (int, np.int32, np.int64)):
symbol = chemical_symbols[symbol]
from hankel import SymmetricFourierTransform
ht = SymmetricFourierTransform(ndim=3, h=self._integration_step, N=self._nodes)
k = np.linspace(0.001, 12, 2000)
n = self.charge(symbol, charge)
# we do not calculate the Fourier Transform at k=0 for numerical reason
fx = ht.transform(n, k[1:] * 2 * np.pi)[0]
# instead it is replaced by the exact value fx(0)=Z
fx = np.concatenate([[atomic_numbers[symbol]], fx])
fx = interp1d(
k, fx, fill_value=(atomic_numbers[symbol], 0.0), bounds_error=False
)
return fx
def _to_lobato(self, symbol, charge):
fx = self.x_ray_scattering_factor(symbol, charge)
n = self.charge(symbol, charge)
r = np.linspace(0, 20, 10000)
Z = np.trapezoid(4 * np.pi * r**2 * n(r), dx=np.diff(r))
fe0 = np.trapezoid(4 * np.pi * r**4 * n(r), dx=np.diff(r)) / units.Bohr / 3
k = np.linspace(0.0, 12, 100)
fe = np.zeros(len(k))
fe[k != 0] = (Z - fx(k[k > 0])) / k[k > 0] ** 2 / (2 * np.pi**2 * units.Bohr)
fe[k == 0] = fe0
if isinstance(symbol, str):
symbol = atomic_numbers[symbol]
parametrization_aeatom = LobatoParametrization({})
parametrization_aeatom.fit(symbol, k, fe)
return parametrization_aeatom
def potential(self, symbol: str, charge: float = 0.0):
return self._to_lobato(symbol, charge).potential(symbol, charge)
def scattering_factor(self, symbol, charge: float = 0.0):
return self._to_lobato(symbol, charge).scattering_factor(symbol, charge)
# def potential(self, symbol: str, charge: float = 0.0):
# return interp1d(k, fe, fill_value=(fe0, 0.), bounds_error=False)
# def potential(self, symbol: str, charge: float = 0.0):
# """
# Calculate the radial electrostatic potential for an atom.
#
# Parameters
# ----------
# symbol : str
# Chemical symbol of the atomic element.
# charge : float, optional
# Charge the atom by the given fractional number of electrons.
#
# Returns
# -------
# potential : callable
# Function of the radial electrostatic potential with parameter 'r'
# corresponding to the radial distance from the core.
# """
#
# ae = self._get_all_electron_atom(symbol, charge)
# r = ae.rgd.r_g * units.Bohr
#
# ve = -ae.rgd.poisson(ae.n_sg.sum(0))
# ve = interp1d(r, ve, fill_value="extrapolate", bounds_error=False)
#
# vr = (
# lambda r: atomic_numbers[symbol] / r / (4 * np.pi * eps0)
# + ve(r) / r * units.Hartree * units.Bohr
# )
# return vr
