Coverage for gpaw/new/density.py: 94%
231 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1from __future__ import annotations
3from math import pi, sqrt
5import numpy as np
6from ase.units import Bohr, Ha
7from gpaw.core.atom_arrays import AtomArrays, AtomDistribution
8from gpaw.core.atom_centered_functions import (AtomArraysLayout,
9 AtomCenteredFunctions)
10from gpaw.core.plane_waves import PWDesc
11from gpaw.core.uniform_grid import UGArray, UGDesc
12from gpaw.gpu import as_np
13from gpaw.mpi import MPIComm
14from gpaw.new import trace, zips
15from gpaw.typing import Array3D, Vector
16from gpaw.utilities import unpack_hermitian, unpack_density
17from gpaw.new.symmetry import SymmetrizationPlan, GPUSymmetrizationPlan
18from gpaw.new.ibzwfs import IBZWaveFunctions
19from gpaw.setup import Setups
22class Density:
23 @classmethod
24 def from_data_and_setups(cls,
25 nt_sR: UGArray,
26 taut_sR: UGArray,
27 D_asii: AtomArrays,
28 charge: float,
29 setups: Setups,
30 nct_aX: AtomCenteredFunctions,
31 tauct_aX: AtomCenteredFunctions) -> Density:
32 xp = nt_sR.xp
33 return cls(nt_sR,
34 taut_sR,
35 D_asii,
36 charge,
37 setups.nvalence + setups.core_charge,
38 [xp.asarray(setup.Delta_iiL) for setup in setups],
39 [setup.Delta0 for setup in setups],
40 [unpack_hermitian(setup.N0_p) for setup in setups],
41 [setup.n_j for setup in setups],
42 [setup.l_j for setup in setups],
43 nct_aX,
44 tauct_aX)
46 @classmethod
47 def from_superposition(cls,
48 *,
49 grid,
50 nct_aX,
51 tauct_aX,
52 atomdist,
53 setups,
54 basis_set,
55 magmom_av,
56 ncomponents,
57 charge=0.0,
58 hund=False,
59 mgga=False):
60 nt_sR = grid.zeros(ncomponents)
61 atom_array_layout = AtomArraysLayout(
62 [(setup.ni, setup.ni) for setup in setups],
63 atomdist=atomdist, dtype=float if ncomponents < 4 else complex)
64 D_asii = atom_array_layout.empty(ncomponents)
65 f_asi = {a: atomic_occupation_numbers(setup,
66 magmom_v,
67 ncomponents,
68 hund,
69 charge / len(setups))
70 for a, (setup, magmom_v)
71 in enumerate(zips(setups, magmom_av))}
72 basis_set.add_to_density(nt_sR.data, f_asi)
73 for a, D_sii in D_asii.items():
74 D_sii[:] = unpack_density(
75 setups[a].initialize_density_matrix(f_asi[a]))
77 xp = nct_aX.xp
78 nt_sR = nt_sR.to_xp(xp)
79 density = cls.from_data_and_setups(nt_sR,
80 None,
81 D_asii.to_xp(xp),
82 charge,
83 setups,
84 nct_aX,
85 tauct_aX)
86 ndensities = ncomponents % 3
87 density.nt_sR.data[:ndensities] += density.nct_R.data
88 if mgga:
89 density.taut_sR = nt_sR.new()
90 density.taut_sR.data[:] = density.tauct_R.data
91 return density
93 def __init__(self,
94 nt_sR: UGArray,
95 taut_sR: UGArray | None,
96 D_asii: AtomArrays,
97 charge: float,
98 nvalence: float,
99 delta_aiiL: list[Array3D],
100 delta0_a: list[float],
101 N0_aii,
102 n_aj: list[list[int]],
103 l_aj: list[list[int]],
104 nct_aX: AtomCenteredFunctions,
105 tauct_aX: AtomCenteredFunctions):
106 self.nt_sR = nt_sR
107 self.taut_sR = taut_sR
108 self.D_asii = D_asii
109 self.delta_aiiL = delta_aiiL
110 self.delta0_a = delta0_a
111 self.N0_aii = N0_aii
112 self.n_aj = n_aj
113 self.l_aj = l_aj
114 self.charge = charge
115 self.nvalence = nvalence
116 self.nct_aX = nct_aX
117 self.tauct_aX = tauct_aX
119 self.grid = nt_sR.desc
120 self.ncomponents = nt_sR.dims[0]
121 self.ndensities = self.ncomponents % 3
122 self.collinear = self.ncomponents != 4
123 self.natoms = len(delta0_a)
125 self._nct_R = None
126 self._tauct_R = None
128 self.symplan = None
130 def __repr__(self):
131 return f'Density({self.nt_sR}, {self.D_asii}, charge={self.charge})'
133 def __str__(self) -> str:
134 return (f'density:\n'
135 f' valence electrons: {self.nvalence}\n'
136 f' components: {self.ncomponents}\n'
137 f' grid points: {self.nt_sR.desc.size}\n'
138 f' charge: {self.charge} # |e|\n')
140 @property
141 def nct_R(self):
142 if self._nct_R is None:
143 self._nct_R = self.grid.empty(xp=self.nt_sR.xp)
144 self.nct_aX.to_uniform_grid(out=self._nct_R,
145 scale=1.0 / (self.ncomponents % 3))
146 return self._nct_R
148 @property
149 def tauct_R(self):
150 if self._tauct_R is None:
151 self._tauct_R = self.grid.empty(xp=self.nt_sR.xp)
152 self.tauct_aX.to_uniform_grid(out=self._tauct_R,
153 scale=1.0 / (self.ncomponents % 3))
154 return self._tauct_R
156 def new(self, new_grid, pw, relpos_ac, atomdist):
157 self.move(relpos_ac, atomdist)
158 new_pw = PWDesc(ecut=0.99 * new_grid.ekin_max(),
159 cell=new_grid.cell,
160 comm=new_grid.comm)
161 old_grid = self.nt_sR.desc
162 old_pw = PWDesc(ecut=0.99 * old_grid.ekin_max(),
163 cell=old_grid.cell,
164 comm=new_grid.comm)
165 new_nt_sR = new_grid.empty(self.ncomponents, xp=self.nt_sR.xp)
166 for new_nt_R, old_nt_R in zips(new_nt_sR, self.nt_sR):
167 old_nt_R.fft(pw=old_pw).morph(new_pw).ifft(out=new_nt_R)
169 self.nct_aX.change_cell(pw)
170 self.tauct_aX.change_cell(pw)
172 return Density(
173 new_nt_sR,
174 None if self.taut_sR is None else new_nt_sR.new(zeroed=True),
175 self.D_asii,
176 self.charge,
177 self.nvalence,
178 self.delta_aiiL,
179 self.delta0_a,
180 self.N0_aii,
181 self.n_aj,
182 self.l_aj,
183 self.nct_aX,
184 self.tauct_aX)
186 @trace
187 def calculate_compensation_charge_coefficients(self) -> AtomArrays:
188 xp = self.D_asii.layout.xp
189 ccc_aL = AtomArraysLayout(
190 [delta_iiL.shape[2] for delta_iiL in self.delta_aiiL],
191 atomdist=self.D_asii.layout.atomdist,
192 xp=xp).empty()
194 for a, D_sii in self.D_asii.items():
195 Q_L = xp.einsum('sij, ijL -> L',
196 D_sii[:self.ndensities].real, self.delta_aiiL[a])
197 Q_L[0] += self.delta0_a[a]
198 ccc_aL[a] = Q_L
200 return ccc_aL
202 def normalize(self, background_charge: float) -> None:
203 comp_charge = 0.0
204 xp = self.D_asii.layout.xp
205 for a, D_sii in self.D_asii.items():
206 comp_charge += xp.einsum('sij, ij ->',
207 D_sii[:self.ndensities].real,
208 self.delta_aiiL[a][:, :, 0])
209 comp_charge += self.delta0_a[a]
210 # comp_charge could be cupy.ndarray:
211 comp_charge = float(comp_charge) * sqrt(4 * pi)
212 comp_charge = self.nt_sR.desc.comm.sum_scalar(comp_charge)
213 charge = comp_charge + self.charge - background_charge
214 pseudo_charge = self.nt_sR[:self.ndensities].integrate().sum()
215 if pseudo_charge != 0.0:
216 x = -charge / pseudo_charge
217 self.nt_sR.data *= x
219 @trace
220 def update(self, ibzwfs: IBZWaveFunctions, ked=False):
221 self.nt_sR.data[:] = 0.0
222 self.D_asii.data[:] = 0.0
223 ibzwfs.add_to_density(self.nt_sR, self.D_asii)
224 self.nt_sR.data[:self.ndensities] += self.nct_R.data
226 # LCAO ...:
227 ibzwfs.normalize_density(self)
229 if ked:
230 self.update_ked(ibzwfs, symmetrize=False)
232 self.symmetrize(ibzwfs.ibz.symmetries)
234 def update_ked(self, ibzwfs, symmetrize=True):
235 if self.taut_sR is None:
236 self.taut_sR = self.nt_sR.new(zeroed=True)
237 else:
238 self.taut_sR.data[:] = 0.0
239 ibzwfs.add_to_ked(self.taut_sR)
240 self.taut_sR.data[:self.ndensities] += self.tauct_R.data
241 if symmetrize:
242 symmetries = ibzwfs.ibz.symmetries
243 self.taut_sR.symmetrize(symmetries.rotation_scc,
244 symmetries.translation_sc)
246 @trace
247 def symmetrize(self, symmetries):
248 self.nt_sR.symmetrize(symmetries.rotation_scc,
249 symmetries.translation_sc)
250 if self.taut_sR is not None:
251 self.taut_sR.symmetrize(symmetries.rotation_scc,
252 symmetries.translation_sc)
254 xp = self.nt_sR.xp
255 if xp is np:
256 D_asii = self.D_asii.gather(broadcast=True, copy=True)
257 if self.symplan is None:
258 self.symplan = SymmetrizationPlan(symmetries, self.l_aj)
259 self.symplan.apply_distributed(D_asii, self.D_asii)
260 else:
261 # GPU version does all the work in rank 0 for now
262 D_asii = self.D_asii.gather(copy=True)
263 if self.D_asii.layout.atomdist.comm.rank == 0:
264 if self.symplan is None:
265 self.symplan = GPUSymmetrizationPlan(
266 symmetries, self.l_aj, D_asii.layout)
267 self.symplan.apply(D_asii.data, D_asii.data)
268 self.D_asii.scatter_from(D_asii)
270 @trace
271 def move(self, relpos_ac, atomdist):
272 self.nt_sR.data[:self.ndensities] -= self.nct_R.data
273 self.nct_aX.move(relpos_ac, atomdist)
274 self.tauct_aX.move(relpos_ac, atomdist)
275 self._nct_R = None
276 self._tauct_R = None
277 self.nt_sR.data[:self.ndensities] += self.nct_R.data
278 self.D_asii = self.D_asii.moved(atomdist)
280 @trace
281 def redist(self,
282 grid: UGDesc,
283 xdesc,
284 atomdist: AtomDistribution,
285 comm1: MPIComm,
286 comm2: MPIComm) -> Density:
287 return Density(
288 self.nt_sR.redist(grid, comm1, comm2),
289 None
290 if self.taut_sR is None else
291 self.taut_sR.redist(grid, comm1, comm2),
292 self.D_asii.redist(atomdist, comm1, comm2),
293 self.charge,
294 self.nvalence,
295 self.delta_aiiL,
296 self.delta0_a,
297 self.N0_aii,
298 self.n_aj,
299 self.l_aj,
300 nct_aX=self.nct_aX.new(xdesc, atomdist),
301 tauct_aX=self.tauct_aX.new(xdesc, atomdist))
303 def calculate_dipole_moment(self, relpos_ac):
304 dip_v = np.zeros(3)
305 ccc_aL = self.calculate_compensation_charge_coefficients()
306 ccc_aL = ccc_aL.to_cpu()
307 pos_av = relpos_ac @ self.nt_sR.desc.cell_cv
308 for a, ccc_L in ccc_aL.items():
309 c = ccc_L[0]
310 dip_v -= c * (4 * pi)**0.5 * pos_av[a]
311 if len(ccc_L) > 1:
312 y, z, x = ccc_L[1:4]
313 dip_v -= np.array([x, y, z]) * (4 * pi / 3)**0.5
314 self.nt_sR.desc.comm.sum(dip_v)
315 for nt_R in self.nt_sR[:self.ndensities]:
316 dip_v -= as_np(nt_R.moment())
317 return dip_v
319 def calculate_orbital_magnetic_moments(self):
320 if self.collinear:
321 from gpaw.new.calculation import CalculationModeError
322 raise CalculationModeError(
323 'Calculator is in collinear mode. '
324 'Collinear calculations require spin–orbit '
325 'coupling for nonzero orbital magnetic moments.')
327 D_asii = self.D_asii
328 if D_asii.layout.size != D_asii.layout.mysize:
329 raise ValueError(
330 'Atomic density matrices should be collected on all '
331 'ranks when calculating orbital magnetic moments.')
333 from gpaw.new.orbmag import calculate_orbmag_from_density
334 return calculate_orbmag_from_density(D_asii, self.n_aj, self.l_aj)
336 def calculate_magnetic_moments(self):
337 magmom_av = np.zeros((self.natoms, 3))
338 magmom_v = np.zeros(3)
339 domain_comm = self.nt_sR.desc.comm
341 if self.ncomponents == 2:
342 for a, D_sii in self.D_asii.items():
343 M_ii = as_np(D_sii[0] - D_sii[1])
344 magmom_av[a, 2] = np.einsum('ij, ij ->', M_ii, self.N0_aii[a])
345 delta_ii = as_np(self.delta_aiiL[a][:, :, 0])
346 magmom_v[2] += (np.einsum('ij, ij ->', M_ii, delta_ii) *
347 sqrt(4 * pi))
348 domain_comm.sum(magmom_av)
349 domain_comm.sum(magmom_v)
351 M_s = self.nt_sR.integrate()
352 magmom_v[2] += M_s[0] - M_s[1]
354 elif self.ncomponents == 4:
355 for a, D_sii in self.D_asii.items():
356 M_vii = D_sii[1:4].real
357 magmom_av[a] = np.einsum('vij, ij -> v',
358 M_vii, self.N0_aii[a])
359 magmom_v += (np.einsum('vij, ij -> v', M_vii,
360 self.delta_aiiL[a][:, :, 0]) *
361 sqrt(4 * pi))
362 domain_comm.sum(magmom_av)
363 domain_comm.sum(magmom_v)
364 magmom_v += self.nt_sR.integrate()[1:]
366 return magmom_v, magmom_av
368 @trace
369 def write_to_gpw(self, writer, flags):
370 D_asp = self.D_asii.to_cpu().to_lower_triangle().gather()
371 nt_sR = self.nt_sR.to_xp(np).gather()
372 if self.taut_sR is not None:
373 taut_sR = self.taut_sR.to_xp(np).gather()
374 if D_asp is None:
375 return # let master do the writing
376 writer.write(
377 density=flags.to_storage_dtype(nt_sR.data * Bohr**-3),
378 atomic_density_matrices=D_asp.data)
379 if self.taut_sR is not None:
380 writer.write(ked=flags.to_storage_dtype(
381 taut_sR.data * (Ha * Bohr**-3)))
384def atomic_occupation_numbers(setup,
385 magmom_v: Vector,
386 ncomponents: int,
387 hund: bool = False,
388 charge: float = 0.0):
389 M = np.linalg.norm(magmom_v)
390 nspins = min(ncomponents, 2)
391 f_si = setup.calculate_initial_occupation_numbers(
392 M, hund, charge=charge, nspins=nspins)
394 if ncomponents == 1:
395 pass
396 elif ncomponents == 2:
397 if magmom_v[2] < 0:
398 f_si = f_si[::-1].copy()
399 else:
400 f_i = f_si.sum(0)
401 fm_i = f_si[0] - f_si[1]
402 f_si = np.zeros((4, len(f_i)))
403 f_si[0] = f_i
404 if M > 0:
405 f_si[1:] = np.asarray(magmom_v)[:, np.newaxis] / M * fm_i
407 return f_si