Coverage for gpaw/core/atom_centered_functions.py: 92%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING, Callable 

4 

5import numpy as np 

6from gpaw.core.atom_arrays import (AtomArrays, AtomArraysLayout, 

7 AtomDistribution) 

8from gpaw.kpt_descriptor import KPointDescriptor 

9from gpaw.lfc import LocalizedFunctionsCollection as LFC 

10from gpaw.mpi import MPIComm, serial_comm 

11from gpaw.new import zips, trace 

12from gpaw.spline import Spline 

13from gpaw.typing import Array1D, ArrayLike2D 

14from gpaw.gpu import XP 

15 

16if TYPE_CHECKING: 

17 from gpaw.core.uniform_grid import UGArray 

18 

19 

20def to_spline(l: int, 

21 rcut: float, 

22 f: Callable[[Array1D], Array1D]) -> Spline: 

23 """Convert to GPAW's Spline object.""" 

24 r = np.linspace(0, rcut, 100) 

25 return Spline.from_data(l, rcut, f(r)) 

26 

27 

28class AtomCenteredFunctions(XP): 

29 def __init__(self, 

30 functions, 

31 relpos_ac: ArrayLike2D, 

32 atomdist: AtomDistribution | None = None, 

33 xp=None): 

34 XP.__init__(self, xp or np) 

35 self.functions = [[to_spline(*f) if isinstance(f, tuple) else f 

36 for f in funcs] 

37 for funcs in functions] 

38 self.relpos_ac = np.array(relpos_ac) 

39 self._atomdist = atomdist 

40 

41 self._layout = None 

42 self._lfc = None 

43 

44 def __repr__(self): 

45 funcs = [['spdfgh'[f.l] for f in ff] for ff in self.functions[:4]] 

46 if len(self.functions) > 4: 

47 funcs.append(...) 

48 return (f'{self.__class__.__name__}' 

49 f'(functions={funcs}, atomdist={self.atomdist})') 

50 

51 def new(self, desc, atomdist): 

52 raise NotImplementedError 

53 

54 @property 

55 def layout(self): 

56 self._lazy_init() 

57 return self._layout 

58 

59 @property 

60 def atomdist(self): 

61 self._lazy_init() 

62 return self._atomdist 

63 

64 def _lazy_init(self): 

65 raise NotImplementedError 

66 

67 def empty(self, 

68 dims: int | tuple[int, ...] = (), 

69 comm: MPIComm = serial_comm) -> AtomArrays: 

70 """Create AtomsArray for coefficients.""" 

71 return self.layout.empty(dims, comm) 

72 

73 def move(self, 

74 relpos_ac: np.ndarray, 

75 atomdist: AtomDistribution) -> AtomDistribution: 

76 """Move atoms to new positions.""" 

77 self.relpos_ac = np.array(relpos_ac) 

78 self._atomdist = atomdist 

79 if self._lfc is not None: 

80 self._layout = self._layout.new(atomdist=atomdist) 

81 migration = self._lfc.set_positions(relpos_ac, atomdist) 

82 if migration: 

83 atomdist = AtomDistribution( 

84 [sphere.rank for sphere in self._lfc.sphere_a], 

85 atomdist.comm) 

86 return atomdist 

87 

88 def add_to(self, functions, coefs=1.0): 

89 """Add atom-centered functions multiplied by *coefs* to *functions*.""" 

90 self._lazy_init() 

91 if isinstance(coefs, float): 

92 self._lfc.add(functions.data, coefs) 

93 else: 

94 self._lfc.add(functions.data, coefs, q=0) 

95 

96 @trace 

97 def integrate(self, functions, out=None, add_to=False): 

98 """Calculate integrals of atom-centered functions multiplied by 

99 *functions*. 

100 """ 

101 self._lazy_init() 

102 if out is None: 

103 assert not add_to 

104 out = self.layout.empty(functions.dims, functions.comm) 

105 self._lfc.integrate(functions.data, out, q=0, add_to=add_to) 

106 return out 

107 

108 def derivative(self, functions, out=None): 

109 """Calculate derivatives of integrals with respect to atom 

110 positions. 

111 """ 

112 self._lazy_init() 

113 if out is None: 

114 out = self.layout.empty(functions.dims + (3,), functions.comm) 

115 coef_axiv = {a: self.xp.moveaxis(array_xvi, -2, -1) 

116 for a, array_xvi in out._arrays.items()} 

117 self._lfc.derivative(functions.data, coef_axiv, q=0) 

118 return out 

119 

120 def stress_contribution(self, a, c=1.0): 

121 self._lazy_init() 

122 return self._lfc.stress_tensor_contribution(a.data, c) 

123 

124 

125class UGAtomCenteredFunctions(AtomCenteredFunctions): 

126 def __init__(self, 

127 functions, 

128 relpos_ac, 

129 grid, 

130 *, 

131 atomdist=None, 

132 integrals=None, 

133 cut=False, 

134 xp=np): 

135 AtomCenteredFunctions.__init__(self, 

136 functions, 

137 relpos_ac, 

138 atomdist, xp=xp) 

139 self.grid = grid 

140 self.integrals = integrals 

141 self.cut = cut 

142 

143 def new(self, grid, atomdist): 

144 return UGAtomCenteredFunctions( 

145 self.functions, 

146 self.relpos_ac, 

147 grid, 

148 atomdist=atomdist, 

149 integrals=self.integrals, 

150 cut=self.cut, 

151 xp=self.xp) 

152 

153 def _lazy_init(self): 

154 if self._lfc is not None: 

155 return 

156 gd = self.grid._gd 

157 kd = KPointDescriptor(np.array([self.grid.kpt])) 

158 self._lfc = LFC(gd, self.functions, kd, 

159 dtype=self.grid.dtype, 

160 integral=self.integrals, 

161 forces=True, 

162 cut=self.cut, 

163 xp=self.xp) 

164 self._lfc.set_positions(self.relpos_ac) 

165 

166 if self._atomdist is None: 

167 self._atomdist = AtomDistribution( 

168 ranks=np.array([sphere.rank for sphere in self._lfc.sphere_a]), 

169 comm=self.grid.comm) 

170 else: 

171 for sphere, rank in zips(self._lfc.sphere_a, 

172 self._atomdist.rank_a): 

173 assert sphere.rank == rank 

174 assert self.grid.comm is self._atomdist.comm 

175 

176 self._layout = AtomArraysLayout([sum(2 * f.l + 1 for f in funcs) 

177 for funcs in self.functions], 

178 self._atomdist, 

179 self.grid.dtype, xp=self.xp) 

180 

181 def to_uniform_grid(self, 

182 out: UGArray, 

183 scale: float = 1.0) -> UGArray: 

184 out.data[:] = 0.0 

185 self.add_to(out, scale) 

186 return out