Source code for gpaw.core.atom_centered_functions

from __future__ import annotations

from typing import TYPE_CHECKING, Callable

import numpy as np
from gpaw.core.atom_arrays import (AtomArrays, AtomArraysLayout,
from gpaw.kpt_descriptor import KPointDescriptor
from gpaw.lfc import LocalizedFunctionsCollection as LFC
from gpaw.mpi import MPIComm, serial_comm
from import zips
from gpaw.spline import Spline
from gpaw.typing import Array1D, ArrayLike2D
from gpaw.gpu import XP

    from gpaw.core.uniform_grid import UGArray

def to_spline(l: int,
              rcut: float,
              f: Callable[[Array1D], Array1D]) -> Spline:
    """Convert to GPAW's Spline object."""
    r = np.linspace(0, rcut, 100)
    return Spline.from_data(l, rcut, f(r))

[docs]class AtomCenteredFunctions(XP): def __init__(self, functions, fracpos_ac: ArrayLike2D, atomdist: AtomDistribution | None = None, xp=None): XP.__init__(self, xp or np) self.functions = [[to_spline(*f) if isinstance(f, tuple) else f for f in funcs] for funcs in functions] self.fracpos_ac = np.array(fracpos_ac) self._atomdist = atomdist self._layout = None self._lfc = None def __repr__(self): funcs = [['spdfgh'[f.l] for f in ff] for ff in self.functions[:4]] if len(self.functions) > 4: funcs.append(...) return (f'{self.__class__.__name__}' f'(functions={funcs}, atomdist={self.atomdist})')
[docs] def new(self, desc, atomdist): raise NotImplementedError
@property def layout(self): self._lazy_init() return self._layout @property def atomdist(self): self._lazy_init() return self._atomdist def _lazy_init(self): raise NotImplementedError
[docs] def empty(self, dims: int | tuple[int, ...] = (), comm: MPIComm = serial_comm) -> AtomArrays: """Create AtomsArray for coefficients.""" return self.layout.empty(dims, comm)
[docs] def move(self, fracpos_ac, atomdist): """Move atoms to new positions.""" self.fracpos_ac = np.array(fracpos_ac) self._atomdist = atomdist if self._lfc is not None: self._layout = self._lfc.set_positions(fracpos_ac, atomdist)
[docs] def add_to(self, functions, coefs=1.0): """Add atom-centered functions multiplied by *coefs* to *functions*.""" self._lazy_init() if isinstance(coefs, float): self._lfc.add(, coefs) else: self._lfc.add(, coefs, q=0)
[docs] def integrate(self, functions, out=None): """Calculate integrals of atom-centered functions multiplied by *functions*. """ self._lazy_init() if out is None: out = self.layout.empty(functions.dims, functions.comm) self._lfc.integrate(, out, q=0) return out
[docs] def derivative(self, functions, out=None): """Calculate derivatives of integrals with respect to atom positions. """ self._lazy_init() if out is None: out = self.layout.empty(functions.dims + (3,), functions.comm) coef_axiv = {a: self.xp.moveaxis(array_xvi, -2, -1) for a, array_xvi in out._arrays.items()} self._lfc.derivative(, coef_axiv, q=0) return out
[docs] def stress_contribution(self, a, c=1.0): self._lazy_init() return self._lfc.stress_tensor_contribution(, c)
class UGAtomCenteredFunctions(AtomCenteredFunctions): def __init__(self, functions, fracpos_ac, grid, *, atomdist=None, integral=None, cut=False, xp=np): AtomCenteredFunctions.__init__(self, functions, fracpos_ac, atomdist, xp=xp) self.grid = grid self.integral = integral self.cut = cut def new(self, grid, atomdist): return UGAtomCenteredFunctions( self.functions, self.fracpos_ac, grid, atomdist=atomdist, integral=self.integral, cut=self.cut, xp=self.xp) def _lazy_init(self): if self._lfc is not None: return gd = self.grid._gd kd = KPointDescriptor(np.array([self.grid.kpt])) self._lfc = LFC(gd, self.functions, kd, dtype=self.grid.dtype, integral=self.integral, forces=True, cut=self.cut, xp=self.xp) self._lfc.set_positions(self.fracpos_ac) if self._atomdist is None: self._atomdist = AtomDistribution( ranks=np.array([sphere.rank for sphere in self._lfc.sphere_a]), comm=self.grid.comm) else: for sphere, rank in zips(self._lfc.sphere_a, self._atomdist.rank_a): assert sphere.rank == rank assert self.grid.comm is self._atomdist.comm self._layout = AtomArraysLayout([sum(2 * f.l + 1 for f in funcs) for funcs in self.functions], self._atomdist, self.grid.dtype, xp=self.xp) def to_uniform_grid(self, out: UGArray, scale: float = 1.0) -> UGArray:[:] = 0.0 self.add_to(out, scale) return out