Coverage for gpaw/pair_density.py: 49%
152 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from math import sqrt, pi
2import numpy as np
4from gpaw.utilities import pack_density, unpack_density
5from gpaw.utilities.tools import pick
6from gpaw.lfc import LocalizedFunctionsCollection as LFC, BasisFunctions
9# XXX Document what is the difference between PairDensity2 and 1.
10class PairDensity2:
11 def __init__(self, density, spos_ac, finegrid):
12 """Initialization needs a paw instance, and whether the compensated
13 pair density should be on the fine grid (boolean)"""
15 self.density = density
16 self.finegrid = finegrid
18 if not finegrid:
19 density.Ghat = LFC(density.gd,
20 [setup.ghat_l
21 for setup in density.setups],
22 integral=sqrt(4 * pi))
23 density.Ghat.set_positions(spos_ac)
25 def initialize(self, kpt, n1, n2):
26 """Set wave function indices."""
27 self.n1 = n1
28 self.n2 = n2
29 self.spin = kpt.s
30 self.P_ani = kpt.P_ani
31 self.psit1_G = pick(kpt.psit_nG, n1)
32 self.psit2_G = pick(kpt.psit_nG, n2)
34 def get_coarse(self, nt_G):
35 """Get pair density"""
36 np.multiply(self.psit1_G.conj(), self.psit2_G, nt_G)
38 def add_compensation_charges(self, nt_G, rhot_g):
39 """Add compensation charges to input pair density, which
40 is interpolated to the fine grid if needed."""
42 if self.finegrid:
43 # interpolate the pair density to the fine grid
44 self.density.interpolator.apply(nt_G, rhot_g)
45 else:
46 # copy values
47 rhot_g[:] = nt_G
49 # Determine the compensation charges for each nucleus
50 Q_aL = {}
51 for a, P_ni in self.P_ani.items():
52 assert P_ni.dtype == float
53 # Generate density matrix
54 P1_i = P_ni[self.n1]
55 P2_i = P_ni[self.n2]
56 D_ii = np.outer(P1_i.conj(), P2_i)
57 # allowed to pack as used in the scalar product with
58 # the symmetric array Delta_pL
59 D_p = pack_density(D_ii, tolerance=1e30)
61 # Determine compensation charge coefficients:
62 Q_aL[a] = np.dot(D_p, self.density.setups[a].Delta_pL)
64 # Add compensation charges
65 if self.finegrid:
66 self.density.ghat.add(rhot_g, Q_aL)
67 else:
68 self.density.Ghat.add(rhot_g, Q_aL)
71class PairDensity:
72 def __init__(self, paw):
73 self.set_paw(paw)
75 def set_paw(self, paw):
76 """basic initialisation knowing the calculator"""
77 self.wfs = paw.wfs
78 self.lcao = paw.wfs.mode == 'lcao'
79 self.density = paw.density
80 self.setups = paw.wfs.setups
81 self.spos_ac = paw.spos_ac
83 self.spin = 0
84 self.k = 0
85 self.weight = 0.0
87 if self.lcao:
88 assert paw.wfs.dtype == float
89 self.wfs = paw.wfs
91 def initialize(self, kpt, i, j):
92 """initialize yourself with the wavefunctions"""
93 self.i = i
94 self.j = j
96 if kpt is not None:
97 self.spin = kpt.s
98 self.k = kpt.k
99 self.weight = kpt.weight
100 self.P_ani = kpt.P_ani
101 if self.lcao:
102 self.q = kpt.q
103 self.wfi_M = kpt.C_nM[i]
104 self.wfj_M = kpt.C_nM[j]
105 else:
106 self.wfi = kpt.psit_nG[i]
107 self.wfj = kpt.psit_nG[j]
109 def get_lcao(self, finegrid=False):
110 """Get pair density"""
111 # Expand the pair density as density matrix
112 rho_MM = (0.5 * np.outer(self.wfi_M, self.wfj_M) +
113 0.5 * np.outer(self.wfj_M, self.wfi_M))
115 rho_G = self.density.gd.zeros()
116 self.wfs.basis_functions.construct_density(rho_MM, rho_G, self.q)
118 if not finegrid:
119 return rho_G
121 # interpolate the pair density to the fine grid
122 rho_g = self.density.finegd.zeros()
123 self.density.interpolator.apply(rho_G, rho_g)
125 return rho_g
127 def get(self, finegrid=False):
128 """Get pair density"""
130 if self.lcao:
131 return self.get_lcao(finegrid)
133 nijt_G = self.wfi.conj() * self.wfj
134 if not finegrid:
135 return nijt_G
137 # interpolate the pair density to the fine grid
138 nijt_g = self.density.finegd.empty(dtype=nijt_G.dtype)
139 self.density.interpolator.apply(nijt_G, nijt_g)
141 return nijt_g
143 def with_compensation_charges(self, finegrid=False):
144 """Get pair density including the compensation charges"""
145 rhot_g = self.get(finegrid)
147 # Determine the compensation charges for each nucleus
148 Q_aL = {}
149 for a, P_ni in self.P_ani.items():
150 # Generate density matrix
151 Pi_i = P_ni[self.i].conj()
152 Pj_i = P_ni[self.j]
153 D_ii = np.outer(Pi_i, Pj_i)
154 # allowed to pack as used in the scalar product with
155 # the symmetric array Delta_pL
156 D_p = pack_density(D_ii)
158 # Determine compensation charge coefficients:
159 Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL)
161 # Add compensation charges
162 if finegrid:
163 self.density.ghat.add(rhot_g, Q_aL)
164 else:
165 if not hasattr(self.density, 'Ghat'):
166 self.density.Ghat = LFC(self.density.gd,
167 [setup.ghat_l
168 for setup in self.setups],
169 integral=sqrt(4 * pi))
170 self.density.Ghat.set_positions(self.spos_ac)
171 self.density.Ghat.add(rhot_g, Q_aL)
173 return rhot_g
175 def with_ae_corrections(self, finegrid=False):
176 """Get pair density including the AE corrections"""
177 nij_g = self.get(finegrid)
179 # Generate the density matrix
180 D_ap = {}
181# D_aii = {}
182 for a, P_ni in self.P_ani.items():
183 Pi_i = P_ni[self.i]
184 Pj_i = P_ni[self.j]
185 D_ii = np.outer(Pi_i.conj(), Pj_i)
186 # Note: D_ii is not symmetric but the products of partial waves are
187 # so that we can pack
188 D_ap[a] = pack_density(D_ii)
189# D_aii[a] = D_ii
191 # Load partial waves if needed
192 if ((finegrid and (not hasattr(self, 'phi'))) or
193 ((not finegrid) and (not hasattr(self, 'Phi')))):
195 # Splines
196 splines = {}
197 phi_aj = []
198 phit_aj = []
199 for a, id in enumerate(self.setups.id_a):
200 if id in splines:
201 phi_j, phit_j = splines[id]
202 else:
203 # Load splines:
204 phi_j, phit_j = self.setups[a].get_partial_waves()[:2]
205 splines[id] = (phi_j, phit_j)
206 phi_aj.append(phi_j)
207 phit_aj.append(phit_j)
209 # Store partial waves as class variables
210 if finegrid:
211 gd = self.density.finegd
212 self.__class__.phi = BasisFunctions(gd, phi_aj)
213 self.__class__.phit = BasisFunctions(gd, phit_aj)
214 self.__class__.phi.set_positions(self.spos_ac)
215 self.__class__.phit.set_positions(self.spos_ac)
216 else:
217 gd = self.density.gd
218 self.__class__.Phi = BasisFunctions(gd, phi_aj)
219 self.__class__.Phit = BasisFunctions(gd, phit_aj)
220 self.__class__.Phi.set_positions(self.spos_ac)
221 self.__class__.Phit.set_positions(self.spos_ac)
223 # Add AE corrections
224 if finegrid:
225 phi = self.phi
226 phit = self.phit
227 gd = self.density.finegd
228 else:
229 phi = self.Phi
230 phit = self.Phit
231 gd = self.density.gd
233 rho_MM = np.zeros((phi.Mmax, phi.Mmax))
234 M1 = 0
235 for a, setup in enumerate(self.setups):
236 ni = setup.ni
237 D_p = D_ap.get(a)
238 if D_p is None:
239 D_p = np.empty(ni * (ni + 1) // 2)
240 if gd.comm.size > 1:
241 gd.comm.broadcast(D_p, self.wfs.partition.rank_a[a])
242 D_ii = unpack_density(D_p)
243# D_ii = D_aii.get(a)
244# if D_ii is None:
245# D_ii = np.empty((ni, ni))
246# if gd.comm.size > 1:
247# gd.comm.broadcast(D_ii, self.wfs.atom_partition.rank_a[a])
248 M2 = M1 + ni
249 rho_MM[M1:M2, M1:M2] = D_ii
250 M1 = M2
252 # construct_density assumes symmetric rho_MM and
253 # takes only the upper half of it
254 phi.construct_density(rho_MM, nij_g, q=-1)
255 phit.construct_density(-rho_MM, nij_g, q=-1)
256 # TODO: use ae_valence_density_correction
257# phi.lfc.ae_valence_density_correction(
258# rho_MM, nij_g, np.zeros(len(phi.M_W), np.intc), np.zeros(self.na))
259# phit.lfc.ae_valence_density_correction(
260# -rho_MM, nij_g, np.zeros(len(phit.M_W), np.intc),
261# np.zeros(self.na))
263 return nij_g