Coverage for gpaw/new/backwards_compatibility.py: 74%
301 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1from functools import cached_property
2from types import SimpleNamespace
4import numpy as np
5from ase import Atoms
6from ase.units import Bohr
8from gpaw.band_descriptor import BandDescriptor
9from gpaw.fftw import MEASURE
10from gpaw.kpt_descriptor import KPointDescriptor
11from gpaw.new import prod, zips
12from gpaw.new.calculation import DFTCalculation
13from gpaw.new.pwfd.wave_functions import PWFDWaveFunctions
14from gpaw.projections import Projections
15from gpaw.pw.descriptor import PWDescriptor
16from gpaw.utilities import pack_density
17from gpaw.utilities.timing import nulltimer
18from gpaw.wavefunctions.arrays import (PlaneWaveExpansionWaveFunctions,
19 UniformGridWaveFunctions)
22class PT:
23 def __init__(self, ibzwfs):
24 self.ibzwfs = ibzwfs
26 def integrate(self, psit_nG, P_ani, q):
27 pt_aiX = self.ibzwfs.wfs_qs[q][0].pt_aiX
28 pt_aiX._lazy_init()
29 pt_aiX._lfc.integrate(psit_nG, P_ani, q=0)
31 def add(self, psit_nG, c_axi, q):
32 self.ibzwfs.wfs_qs[q][0].pt_aiX._lfc.add(psit_nG, c_axi, q=0)
34 def dict(self, shape):
35 return self.ibzwfs.wfs_qs[0][0].pt_aiX.empty(shape,
36 self.ibzwfs.band_comm)
39class FakeWFS:
40 def __init__(self,
41 ibzwfs,
42 density,
43 potential,
44 setups,
45 comm,
46 occ_calc,
47 hamiltonian,
48 atoms: Atoms,
49 scale_pw_coefs=False):
50 from gpaw.utilities.partition import AtomPartition
51 self.timer = nulltimer
52 self.setups = setups
53 self.ibzwfs = ibzwfs
54 self.density = density
55 self.potential = potential
56 self.hamiltonian = hamiltonian
57 ibz = ibzwfs.ibz
58 self.kd = kd = KPointDescriptor(ibz.bz.kpt_Kc, ibzwfs.nspins)
59 kd.ibzk_kc = ibz.kpt_kc
60 kd.weight_k = ibz.weight_k
61 kd.sym_k = ibz.s_K
62 kd.time_reversal_k = ibz.time_reversal_K
63 kd.bz2ibz_k = ibz.bz2ibz_K
64 kd.ibz2bz_k = ibz.ibz2bz_k
65 kd.bz2bz_ks = ibz.bz2bz_Ks
66 kd.nibzkpts = len(ibz)
67 kd.symmetry = ibz.symmetries._old_symmetry
68 kd.set_communicator(ibzwfs.kpt_comm)
69 self.bd = BandDescriptor(ibzwfs.nbands, ibzwfs.band_comm)
70 self.grid = density.nt_sR.desc
71 self.gd = self.grid._gd
72 atomdist = density.D_asii.layout.atomdist
73 self.atom_partition = AtomPartition(atomdist.comm, atomdist.rank_a)
74 # self.setups.set_symmetry(ibzwfs.ibz.symmetries.symmetry)
75 self.occ_calc = occ_calc
76 self.occupations = occ_calc.occ
77 self.nvalence = int(round(density.nvalence))
78 self.nvalence = density.nvalence
79 # assert self.nvalence == density.nvalence
80 self.world = comm
81 if ibzwfs.fermi_levels is not None:
82 self.fermi_levels = ibzwfs.fermi_levels
83 if len(self.fermi_levels) == 1:
84 self.fermi_level = self.fermi_levels[0]
85 self.nspins = ibzwfs.nspins
86 self.dtype = ibzwfs.dtype
87 wfs = ibzwfs.wfs_qs[0][0]
88 self.pd = None
89 self.basis_functions = getattr(wfs, # dft.scf_loop.hamiltonian,
90 'basis', None)
91 if isinstance(wfs, PWFDWaveFunctions):
92 if hasattr(wfs.psit_nX.desc, 'ecut'):
93 self.mode = 'pw'
94 self.ecut = wfs.psit_nX.desc.ecut
95 self.pd = PWDescriptor(self.ecut,
96 self.gd, self.dtype, self.kd, _new=True)
97 self.pwgrid = self.grid.new(dtype=self.dtype)
98 else:
99 self.mode = 'fd'
100 else:
101 self.mode = 'lcao'
102 self.manytci = wfs.tci_derivatives.manytci
103 if self.basis_functions is not None:
104 self.ksl = SimpleNamespace(Mstart=self.basis_functions.Mstart,
105 Mstop=self.basis_functions.Mstop)
106 self.collinear = wfs.ncomponents < 4
107 self.positions_set = True
108 self.read_from_file_init_wfs_dm = ibzwfs.read_from_file_init_wfs_dm
110 self.pt = PT(ibzwfs)
111 self.scalapack_parameters = (None, 1, 1, 128)
112 self.ngpts = prod(self.gd.N_c)
113 if self.mode == 'pw' and scale_pw_coefs:
114 self.scale = self.ngpts
115 else:
116 self.scale = 1
117 self.fftwflags = MEASURE
119 def apply_pseudo_hamiltonian(self, kpt, ham, a1, a2):
120 desc = self.ibzwfs.wfs_qs[kpt.q][0].psit_nX.desc
121 self.hamiltonian.apply(
122 self.potential.vt_sR,
123 None,
124 self.ibzwfs, # needed for hybrids
125 getattr(ham, 'D_asii', None), # needed for hybrids
126 desc.from_data(data=a1),
127 desc.from_data(data=a2),
128 kpt.s)
130 def calculate_occupation_numbers(self, fixed):
131 self.ibzwfs.calculate_occs(
132 self.occ_calc,
133 fix_fermi_level=fixed)
135 def empty(self, n, q):
136 return np.empty((n,) +
137 self.ibzwfs.wfs_qs[q][0].psit_nX.data.shape[1:],
138 complex if self.mode == 'pw' else self.dtype)
140 @cached_property
141 def work_array(self):
142 return np.empty(
143 (self.bd.mynbands,) + self.ibzwfs.get_max_shape(),
144 complex if self.mode == 'pw' else self.dtype)
146 @cached_property
147 def work_matrix_nn(self):
148 from gpaw.matrix import Matrix
149 return Matrix(
150 self.bd.nbands, self.bd.nbands,
151 dtype=self.dtype,
152 dist=(self.bd.comm, self.bd.comm.size))
154 @property
155 def orthonormalized(self):
156 return self.ibzwfs.wfs_qs[0][0].orthonormalized
158 def orthonormalize(self, kpt=None):
159 if kpt is None:
160 kpts = list(self.ibzwfs)
161 else:
162 kpts = [self.ibzwfs.wfs_qs[kpt.q][kpt.s]]
163 for wfs in kpts:
164 wfs._P_ani = None
165 wfs.orthonormalized = False
166 wfs.orthonormalize()
168 def make_preconditioner(self, blocksize):
169 if self.mode == 'pw':
170 from gpaw.wavefunctions.pw import Preconditioner
171 return Preconditioner(self.pd.G2_qG, self.pd,
172 _scale=self.ngpts**2)
173 from gpaw.preconditioner import Preconditioner
174 return Preconditioner(self.gd, self.hamiltonian.kin, self.dtype,
175 blocksize)
177 def _get_wave_function_array(self, u, n, realspace=True, periodic=False):
178 assert realspace and not periodic
179 psit_X = self.kpt_u[u].wfs.psit_nX[n]
180 if hasattr(psit_X, 'ifft'):
181 psit_R = psit_X.ifft(grid=self.pwgrid, periodic=True)
182 psit_R.multiply_by_eikr(psit_X.desc.kpt_c)
183 return psit_R.data
184 return psit_X.data
186 def get_wave_function_array(self, n, k, s,
187 realspace=True,
188 periodic=False,
189 cut=False):
190 assert not cut
191 assert self.ibzwfs.band_comm.size == 1
192 assert self.ibzwfs.kpt_comm.size == 1
193 if self.mode == 'lcao':
194 assert not realspace
195 return self.kpt_qs[k][s].C_nM[n]
196 psit_X = self.kpt_qs[k][s].wfs.psit_nX[n]
197 if not realspace:
198 return psit_X.data
199 if self.mode == 'pw':
200 psit_R = psit_X.ifft(grid=self.pwgrid, periodic=True)
201 if not periodic:
202 psit_R.multiply_by_eikr(psit_X.desc.kpt_c)
203 else:
204 psit_R = psit_X
205 if periodic:
206 psit_R.multiply_by_eikr(-psit_R.desc.kpt_c)
207 return psit_R.data
209 def collect_projections(self, k, s):
210 return self.kpt_qs[k][s].projections.collect()
212 def collect_eigenvalues(self, k, s):
213 return self.ibzwfs.wfs_qs[k][s].eig_n.copy()
215 @cached_property
216 def kpt_u(self):
217 return [kpt
218 for kpt_s in self.kpt_qs
219 for kpt in kpt_s]
221 @cached_property
222 def kpt_qs(self):
223 return [[KPT(self.mode, wfs, self.atom_partition, self.scale,
224 self.pd, self.gd)
225 for wfs in wfs_s]
226 for wfs_s in self.ibzwfs.wfs_qs]
228 def integrate(self, a_nX, b_nX, global_integral):
229 if self.mode == 'fd':
230 return self.gd.integrate(a_nX, b_nX, global_integral)
231 x = self.pd.integrate(a_nX, b_nX, global_integral)
232 return self.ngpts**2 * x
235class KPT:
236 def __init__(self, mode, wfs, atom_partition, scale, pd, gd):
237 self.mode = mode
238 self.scale = scale
239 self.wfs = wfs
240 self.pd = pd
241 self.gd = gd
243 try:
244 I1 = 0
245 nproj_a = []
246 for a, shape in enumerate(wfs.P_ani.layout.shape_a):
247 I2 = I1 + prod(shape)
248 nproj_a.append(I2 - I1)
249 I1 = I2
250 except RuntimeError:
251 pass
252 else:
253 self.projections = Projections(
254 wfs.nbands,
255 nproj_a,
256 atom_partition,
257 wfs.P_ani.comm,
258 wfs.ncomponents < 4,
259 wfs.spin,
260 data=wfs.P_ani.data)
262 self.s = wfs.spin if wfs.ncomponents < 4 else None
263 self.k = wfs.k
264 self.q = wfs.q
265 self.weight = wfs.spin_degeneracy * wfs.weight
266 self.weightk = wfs.weight
267 if isinstance(wfs, PWFDWaveFunctions):
268 self.psit_nX = wfs.psit_nX
269 else:
270 self.C_nM = wfs.C_nM.data
271 self.S_MM = wfs.S_MM.data
272 self.P_aMi = wfs.P_aMi
273 if mode == 'fd':
274 self.phase_cd = wfs.psit_nX.desc.phase_factor_cd
276 @property
277 def P_ani(self):
278 return self.wfs.P_ani
280 @property
281 def eps_n(self):
282 return self.wfs.myeig_n
284 @property
285 def f_n(self):
286 f_n = self.wfs.myocc_n * self.weight
287 f_n.flags.writeable = False
288 return f_n
290 @f_n.setter
291 def f_n(self, val):
292 self.wfs.myocc_n[:] = val / self.weight
294 @property
295 def psit_nG(self):
296 if not hasattr(self, 'psit_nX'):
297 return None
298 data = self.psit_nX.data
299 if self.scale == 1:
300 return data
301 if 1: # isinstance(data, np.ndarray):
302 return data * self.scale
303 data.scale *= self.scale
304 return data
306 @cached_property
307 def psit(self):
308 band_comm = self.psit_nX.comm
309 if self.mode == 'pw':
310 return PlaneWaveExpansionWaveFunctions(
311 self.wfs.nbands, self.pd, self.wfs.dtype,
312 self.psit_nG,
313 kpt=self.q,
314 dist=(band_comm, band_comm.size),
315 spin=self.s,
316 collinear=self.wfs.ncomponents != 4)
317 return UniformGridWaveFunctions(
318 self.wfs.nbands, self.gd, self.wfs.dtype,
319 self.psit_nX.data,
320 kpt=self.q,
321 dist=(band_comm, band_comm.size),
322 spin=self.s,
323 collinear=self.wfs.ncomponents != 4)
326class FakeDensity:
327 def __init__(self, dft: DFTCalculation):
328 self.setups = dft.setups
329 self.D_asii = dft.density.D_asii
330 self.atom_partition = dft._atom_partition
331 try:
332 self.interpolate = dft.pot_calc._interpolate_density
333 self.finegd = dft.pot_calc.fine_grid._gd
334 except AttributeError:
335 pass
336 self.nt_sR = dft.density.nt_sR
337 self.nt_sG = self.nt_sR.data
338 self.gd = self.nt_sR.desc._gd
339 self._densities = dft.densities()
340 self.ncomponents = len(self.nt_sG)
341 self.nspins = self.ncomponents % 3
342 self.collinear = self.ncomponents < 4
344 @cached_property
345 def D_asp(self):
346 D_asp = self.setups.empty_atomic_matrix(self.ncomponents,
347 self.atom_partition)
348 D_asp.update({a: np.array([pack_density(D_ii) for D_ii in D_sii.real])
349 for a, D_sii in self.D_asii.items()})
350 return D_asp
352 @cached_property
353 def nt_sg(self):
354 return self.interpolate(self.nt_sR)[0].data
356 def interpolate_pseudo_density(self):
357 pass
359 def get_all_electron_density(self, *, atoms, gridrefinement):
360 n_sr = self._densities.all_electron_densities(
361 grid_refinement=gridrefinement).scaled(1 / Bohr, Bohr**3)
362 return n_sr.data, n_sr.desc._gd
365class FakeHamiltonian:
366 def __init__(self, ibzwfs, density, potential, pot_calc,
367 e_total_free=np.nan,
368 e_xc=np.nan):
369 self.pot_calc = pot_calc
370 self.ibzwfs = ibzwfs
371 self.density = density
372 self.potential = potential
373 try:
374 self.finegd = self.pot_calc.fine_grid._gd
375 except AttributeError:
376 pass
377 self.grid = potential.vt_sR.desc
378 self.e_total_free = e_total_free
379 self.e_xc = e_xc
381 def update(self, dens, wfs, kin_en_using_band=True):
382 self.potential, _ = self.pot_calc.calculate(
383 self.density, self.ibzwfs, self.potential.vHt_x)
385 energies = self.potential.energies
386 self.e_xc = energies['xc']
387 self.e_coulomb = energies['coulomb']
388 self.e_zero = energies['zero']
389 self.e_external = energies['external']
391 if kin_en_using_band:
392 self.e_kinetic0 = energies['kinetic']
393 else:
394 self.e_kinetic0 = self.ibzwfs.calculate_kinetic_energy(
395 wfs.hamiltonian, self.density)
396 self.ibzwfs.energies['exx_kinetic'] = 0.0
397 energies['kinetic'] = self.e_kinetic0
399 def get_energy(self, e_entropy, wfs, kin_en_using_band=True, e_sic=None):
400 self.e_band = self.ibzwfs.energies['band']
401 if kin_en_using_band:
402 self.e_kinetic = self.e_kinetic0 + self.e_band
403 else:
404 self.e_kinetic = self.e_kinetic0
405 self.e_entropy = e_entropy
406 if 0:
407 print(self.e_kinetic0,
408 self.e_band,
409 self.e_coulomb,
410 self.e_external,
411 self.e_zero,
412 self.e_xc,
413 self.e_entropy)
414 self.e_total_free = (self.e_kinetic + self.e_coulomb +
415 self.e_external + self.e_zero + self.e_xc +
416 self.e_entropy)
418 if e_sic is not None:
419 self.e_sic = e_sic
420 self.e_total_free += e_sic
422 self.e_total_extrapolated = (
423 self.e_total_free +
424 self.ibzwfs.energies['extrapolation'])
426 return self.e_total_free
428 def restrict_and_collect(self, vxct_sg):
429 fine_grid = self.pot_calc.fine_grid
430 vxct_sr = fine_grid.empty(len(vxct_sg))
431 vxct_sr.data[:] = vxct_sg
432 vxct_sR = self.grid.empty(len(vxct_sg))
433 for vxct_r, vxct_R in zips(vxct_sr, vxct_sR):
434 self.pot_calc.restrict(vxct_r, vxct_R)
435 return vxct_sR.data
437 @property
438 def xc(self):
439 return self.pot_calc.xc.xc
441 def dH(self, P, out):
442 for a, I1, I2 in P.indices:
443 dH_ii = self.potential.dH_asii[a][P.spin]
444 out.array[:, I1:I2] = np.dot(P.array[:, I1:I2], dH_ii)