Coverage for gpaw/densities.py: 51%
136 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from __future__ import annotations
3from math import pi
4from typing import TYPE_CHECKING
6import numpy as np
7from ase.units import Bohr
9from gpaw.core.atom_arrays import AtomArrays, AtomArraysLayout
10from gpaw.core.uniform_grid import UGArray
11from gpaw.setup import Setups
12from gpaw.spherical_harmonics import Y
13from gpaw.spline import Spline
14from gpaw.typing import Array1D, Array3D, Vector, Array2D
15from gpaw.new import zips as zip
17if TYPE_CHECKING:
18 from gpaw.new.calculation import DFTCalculation
21class Densities:
22 def __init__(self,
23 nt_sR: UGArray,
24 D_asii: AtomArrays,
25 relpos_ac: Array2D,
26 setups: Setups):
27 self.nt_sR = nt_sR
28 self.D_asii = D_asii
29 self.relpos_ac = relpos_ac
30 self.setups = setups
32 @classmethod
33 def from_calculation(cls, calculation: DFTCalculation):
34 density = calculation.density
35 return cls(density.nt_sR,
36 density.D_asii,
37 calculation.relpos_ac,
38 calculation.setups)
40 def pseudo_densities(self,
41 grid_spacing: float = None, # Ang
42 grid_refinement: int = None,
43 add_compensation_charges: bool = True
44 ) -> UGArray:
45 nt_sR = self._pseudo_densities(grid_spacing, grid_refinement)
47 ncomponents = nt_sR.dims[0]
48 ndensities = ncomponents % 3
50 if add_compensation_charges:
51 cc_asL = AtomArraysLayout(
52 [(ncomponents, setup.Delta_iiL.shape[2])
53 for setup in self.setups],
54 atomdist=self.D_asii.layout.atomdist).empty()
56 for a, D_sii in self.D_asii.items():
57 Q_sL = np.einsum('sij, ijL -> sL',
58 D_sii.real, self.setups[a].Delta_iiL)
59 delta = (self.setups[a].Delta0 +
60 self.setups[a].Nv / (4 * pi)**0.5)
61 Q_sL[:ndensities, 0] += delta / ndensities
62 cc_asL[a] = Q_sL
64 ghat_aLR = self.setups.create_compensation_charges(
65 nt_sR.desc,
66 self.relpos_ac,
67 self.D_asii.layout.atomdist)
68 ghat_aLR.add_to(nt_sR, cc_asL)
70 return nt_sR.scaled(Bohr, Bohr**-3)
72 def _pseudo_densities(self,
73 grid_spacing: float = None, # Ang
74 grid_refinement: int = None,
75 ) -> UGArray:
76 nt_sR = self.nt_sR.to_pbc_grid()
77 grid = nt_sR.desc
78 if grid_spacing is not None:
79 assert grid_refinement is None
80 grid = grid.uniform_grid_with_grid_spacing(
81 grid_spacing / Bohr)
82 elif grid_refinement is not None and grid_refinement > 1:
83 grid = grid.new(size=grid.size * grid_refinement)
84 else:
85 return nt_sR.copy()
87 return nt_sR.interpolate(grid=grid)
89 def all_electron_densities(self,
90 *,
91 grid_spacing: float = None, # Ang
92 grid_refinement: int = None,
93 skip_core: bool = False,
94 ) -> UGArray:
95 n_sR = self._pseudo_densities(grid_spacing, grid_refinement)
96 ncomponents = n_sR.dims[0]
97 nspins = ncomponents % 3
98 grid = n_sR.desc
100 electrons_as = np.zeros((len(self.relpos_ac), ncomponents))
101 splines = {}
102 for a, D_sii in self.D_asii.items():
103 D_sii = D_sii.real
104 relpos_c = self.relpos_ac[a]
105 setup = self.setups[a]
106 if setup not in splines:
107 phi_j, phit_j, nc, nct = setup.get_partial_waves()[:4]
108 if skip_core:
109 nc = Spline.from_data(0, 10.0, [0.0, 0.0])
110 rcut = max(setup.rcut_j)
111 splines[setup] = (rcut, phi_j, phit_j, nc, nct)
112 rcut, phi_j, phit_j, nc, nct = splines[setup]
114 # Expected integral of PAW correction:
115 electrons_s = np.zeros(ncomponents)
116 if skip_core:
117 electrons_s[:nspins] = -setup.Nct / nspins
118 else:
119 electrons_s[:nspins] = (setup.Nc - setup.Nct) / nspins
120 electrons_s += (4 * pi)**0.5 * np.einsum('sij, ij -> s',
121 D_sii,
122 setup.Delta_iiL[:, :, 0])
124 # Add PAW correction:
125 R_v = relpos_c @ grid.cell_cv
126 electrons_s -= add(R_v, n_sR, phi_j, phit_j, nc, nct, rcut, D_sii)
127 electrons_as[a] = electrons_s
129 if not skip_core:
130 # Add missing charge to grid-points closest to atoms:
131 grid.comm.sum(electrons_as)
132 R_ac = np.around(grid.size * self.relpos_ac).astype(int)
133 R_ac %= grid.size
134 for R_c, electrons_s in zip(R_ac, electrons_as):
135 R_c -= grid.start_c
136 if (R_c >= 0).all() and (R_c < grid.mysize_c).all():
137 for n_R, e in zip(n_sR.data, electrons_s):
138 n_R[tuple(R_c)] += e / grid.dv
140 return n_sR.scaled(Bohr, Bohr**-3)
142 def spin_contamination(self, majority_spin=None):
143 """Calculate the spin contamination.
145 Spin contamination is defined as the integral over the
146 spin density difference, where it is negative (i.e. the
147 minority spin density is larger than the majority spin density.
148 """
149 n_sR = self.all_electron_densities()
150 m0, m1 = n_sR.integrate()
151 if majority_spin is None:
152 majority_spin = int(m1 > m0)
153 d_R = n_sR[0].data - n_sR[1].data
154 if majority_spin == 0:
155 d_R *= -1.0
156 d_R = np.where(d_R > 0, d_R, 0.0)
157 return n_sR.desc.from_data(d_R).integrate()
160def add(R_v: Vector,
161 a_sR: UGArray,
162 phi_j: list[Spline],
163 phit_j: list[Spline],
164 nc: Spline,
165 nct: Spline,
166 rcut: float,
167 D_sii: Array3D) -> Array1D:
168 """Add PAW corrections to real-space grid.
170 Returns number of elctrons added.
171 """
172 ug = a_sR.desc
173 R_Rv = ug.xyz()
174 lmax = max(phi.l for phi in phi_j)
175 ncomponents = a_sR.dims[0]
176 nspins = ncomponents % 3
177 electrons_s = np.zeros(ncomponents)
178 start_c = 0 - ug.pbc
179 stop_c = 1 + ug.pbc
180 for u0 in range(start_c[0], stop_c[0]):
181 for u1 in range(start_c[1], stop_c[1]):
182 for u2 in range(start_c[2], stop_c[2]):
183 d_Rv = R_Rv - (R_v + (u0, u1, u2) @ ug.cell_cv)
184 d_R = (d_Rv**2).sum(3)**0.5
185 mask_R = d_R < rcut
186 npoints = mask_R.sum()
187 if npoints == 0:
188 continue
190 a_sr = np.zeros((ncomponents, npoints))
191 d_rv = d_Rv[mask_R]
192 d_r = d_R[mask_R]
193 Y_Lr = [Y(L, *d_rv.T) for L in range((lmax + 1)**2)]
194 phi_jr = [phi.map(d_r) for phi in phi_j]
195 phit_jr = [phit.map(d_r) for phit in phit_j]
196 l_j = [phi.l for phi in phi_j]
198 i1 = 0
199 for l1, phi1_r, phit1_r in zip(l_j, phi_jr, phit_jr):
200 i2 = 0
201 i1b = i1 + 2 * l1 + 1
202 D_smi = D_sii[:, i1:i1b]
203 for l2, phi2_r, phit2_r in zip(l_j, phi_jr, phit_jr):
204 i2b = i2 + 2 * l2 + 1
205 D_smm = D_smi[:, :, i2:i2b]
206 b_sr = np.einsum(
207 'smn, mr, nr -> sr',
208 D_smm,
209 Y_Lr[l1**2:(l1 + 1)**2],
210 Y_Lr[l2**2:(l2 + 1)**2]) * (
211 phi1_r * phi2_r - phit1_r * phit2_r)
212 a_sr += b_sr
213 i2 = i2b
214 i1 = i1b
216 dn_r = nc.map(d_r) - nct.map(d_r)
217 a_sr[:nspins] += dn_r * ((4 * pi)**-0.5 / nspins)
218 electrons_s += a_sr.sum(1) * a_sR.desc.dv
219 a_sR.data[:, mask_R] += a_sr
220 return electrons_s