Coverage for gpaw/lcao/overlap.py: 87%
458 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1"""Module for evaluating two-center integrals.
3Contains classes for evaluating integrals of the form::
5 /
6 | _ _a _ _b _
7 Theta = | f(r - R ) g(r - R ) dr ,
8 |
9 /
11with f and g each being given as a radial function times a spherical
12harmonic.
14Important classes
15-----------------
17Low-level:
19 * OverlapExpansion: evaluate the overlap between a pair of functions (or a
20 function with itself) for some displacement vector: <f | g>. An overlap
21 expansion can be created once for a pair of splines f and g, and actual
22 values of the overlap can then be evaluated for several different
23 displacement vectors.
24 * FourierTransformer: create OverlapExpansion object from pair of splines.
26Mid-level:
28 * TwoSiteOverlapExpansions: evaluate overlaps between two *sets* of functions,
29 where functions in the same set reside on the same location: <f_j | g_j>.
30 * TwoSiteOverlapCalculator: create TwoSiteOverlapExpansions object from
31 pair of lists of splines.
33High-level:
35 * ManySiteOverlapExpansions: evaluate overlaps with many functions in many
36 locations: <f_aj | g_aj>.
37 * ManySiteOverlapCalculator: create ManySiteOverlapExpansions object from
38 pair of lists of splines nested by atom and orbital number.
40The low-level classes do the actual work, while the higher-level ones depend
41on the lower-level ones.
43"""
45from math import pi
47import numpy as np
49from ase.neighborlist import PrimitiveNeighborList
50from ase.data import covalent_radii
51from ase.units import Bohr
53import gpaw.cgpaw as cgpaw
54from gpaw.ffbt import ffbt, FourierBesselTransformer
55from gpaw.gaunt import gaunt
56from gpaw.spherical_harmonics import Yl, nablarlYL
57from gpaw.spline import Spline
58from gpaw.utilities.tools import tri2full
59from gpaw.utilities.timing import nulltimer
61timer = nulltimer # XXX global timer object, only for hacking
63UL = 'L'
66class BaseOverlapExpansionSet:
67 def __init__(self, shape):
68 self.shape = shape
70 def zeros(self, shape=(), dtype=float):
71 return np.zeros(shape + self.shape, dtype=dtype)
74class OverlapExpansion(BaseOverlapExpansionSet):
75 """A list of real-space splines representing an overlap integral."""
76 def __init__(self, la, lb, spline_l):
77 self.la = la
78 self.lb = lb
79 self.lmaxgaunt = max(la, lb)
80 self.spline_l = spline_l
81 self.lmaxspline = (la + lb) % 2 + 2 * len(self.spline_l)
82 BaseOverlapExpansionSet.__init__(self, (2 * la + 1, 2 * lb + 1))
83 self.cspline_l = [spline.spline for spline in self.spline_l]
85 def evaluate(self, r, rlY_lm, G_LLL, x_mi, _nil=np.empty(0)):
86 cgpaw.tci_overlap(self.la, self.lb, G_LLL, self.cspline_l,
87 r, rlY_lm, x_mi,
88 False, _nil, _nil, _nil)
90 def derivative(self, r, Rhat_c, rlY_L, G_LLL, drlYdR_Lc, dxdR_cmi,
91 _nil=np.empty(0)):
92 # timer.start('deriv')
93 cgpaw.tci_overlap(self.la, self.lb, G_LLL, self.cspline_l,
94 r, rlY_L, _nil,
95 True, Rhat_c, drlYdR_Lc, dxdR_cmi)
96 # timer.stop('deriv')
99class TwoSiteOverlapExpansions(BaseOverlapExpansionSet):
100 def __init__(self, la_j, lb_j, oe_jj):
101 self.oe_jj = oe_jj
102 shape = (sum([2 * l + 1 for l in la_j]),
103 sum([2 * l + 1 for l in lb_j]))
104 BaseOverlapExpansionSet.__init__(self, shape)
105 if oe_jj.size == 0:
106 self.lmaxgaunt = 0
107 self.lmaxspline = 0
108 else:
109 self.lmaxgaunt = max(oe.lmaxgaunt for oe in oe_jj.ravel())
110 self.lmaxspline = max(oe.lmaxspline for oe in oe_jj.ravel())
112 def slice(self, x_xMM):
113 assert x_xMM.shape[-2:] == self.shape
114 Ma1 = 0
115 for j1, oe_j in enumerate(self.oe_jj):
116 Mb1 = 0
117 Ma2 = Ma1
118 for j2, oe in enumerate(oe_j):
119 Ma2 = Ma1 + oe.shape[0]
120 Mb2 = Mb1 + oe.shape[1]
121 yield x_xMM[..., Ma1:Ma2, Mb1:Mb2], oe
122 Mb1 = Mb2
123 Ma1 = Ma2
125 def evaluate(self, r, rlY_lm):
126 timer.start('tsoe eval')
127 x_MM = self.zeros()
128 G_LLL = gaunt(self.lmaxgaunt)
129 rlY_L = rlY_lm.toarray(self.lmaxspline)
130 for x_mm, oe in self.slice(x_MM):
131 oe.evaluate(r, rlY_L, G_LLL, x_mm)
132 timer.stop('tsoe eval')
133 return x_MM
135 def derivative(self, r, Rhat, rlY_lm, drlYdR_lmc):
136 x_cMM = self.zeros((3,))
137 G_LLL = gaunt(self.lmaxgaunt)
138 rlY_L = rlY_lm.toarray(self.lmaxspline)
139 drlYdR_Lc = drlYdR_lmc.toarray(self.lmaxspline)
140 # print(drlYdR_lmc.shape)
141 for x_cmm, oe in self.slice(x_cMM):
142 oe.derivative(r, Rhat, rlY_L, G_LLL, drlYdR_Lc, x_cmm)
143 return x_cMM
146class ManySiteOverlapExpansions(BaseOverlapExpansionSet):
147 def __init__(self, tsoe_II, I1_a, I2_a):
148 self.tsoe_II = tsoe_II
149 self.I1_a = I1_a
150 self.I2_a = I2_a
152 M1 = 0
153 M1_a = []
154 for I in I1_a:
155 M1_a.append(M1)
156 M1 += tsoe_II[I, 0].shape[0]
157 self.M1_a = M1_a
159 M2 = 0
160 M2_a = []
161 for I in I2_a:
162 M2_a.append(M2)
163 M2 += tsoe_II[0, I].shape[1]
164 self.M2_a = M2_a
166 shape = (sum([tsoe_II[I, 0].shape[0] for I in I1_a]),
167 sum([tsoe_II[0, I].shape[1] for I in I2_a]))
168 assert (M1, M2) == shape
169 BaseOverlapExpansionSet.__init__(self, shape)
171 def get(self, a1, a2):
172 return self.tsoe_II[self.I1_a[a1], self.I2_a[a2]]
174 def getslice(self, a1, a2, x_xMM):
175 I1 = self.I1_a[a1]
176 I2 = self.I2_a[a2]
177 tsoe = self.tsoe_II[I1, I2]
178 Mstart1 = self.M1_a[a1]
179 Mend1 = Mstart1 + tsoe.shape[0]
180 Mstart2 = self.M2_a[a2]
181 Mend2 = Mstart2 + tsoe.shape[1]
182 return x_xMM[..., Mstart1:Mend1, Mstart2:Mend2], tsoe
184 def evaluate_slice(self, disp, x_qxMM):
185 x_qxmm, oe = self.getslice(disp.a1, disp.a2, x_qxMM)
186 disp.evaluate_overlap(oe, x_qxmm)
189class DomainDecomposedExpansions(BaseOverlapExpansionSet):
190 def __init__(self, msoe, local_indices):
191 self.msoe = msoe
192 self.local_indices = local_indices
193 BaseOverlapExpansionSet.__init__(self, msoe.shape)
195 def evaluate_slice(self, disp, x_xqMM):
196 if disp.a2 in self.local_indices:
197 self.msoe.evaluate_slice(disp, x_xqMM)
200class ManySiteDictionaryWrapper(DomainDecomposedExpansions):
201 # Used with dictionaries such as P_aqMi and dPdR_aqcMi
203 def getslice(self, a1, a2, xdict_aqxMi):
204 msoe = self.msoe
205 tsoe = msoe.tsoe_II[msoe.I1_a[a1], msoe.I2_a[a2]]
206 Mstart = self.msoe.M1_a[a1]
207 Mend = Mstart + tsoe.shape[0]
208 return xdict_aqxMi[a2][..., Mstart:Mend, :], tsoe
210 def evaluate_slice(self, disp, x_aqxMi):
211 if disp.a2 in x_aqxMi:
212 x_qxmi, oe = self.getslice(disp.a1, disp.a2, x_aqxMi)
213 disp.evaluate_overlap(oe, x_qxmi)
214 if disp.a1 in x_aqxMi and (disp.a1 != disp.a2):
215 x2_qxmi, oe2 = self.getslice(disp.a2, disp.a1, x_aqxMi)
216 rdisp = disp.reverse()
217 rdisp.evaluate_overlap(oe2, x2_qxmi)
220class BlacsOverlapExpansions(BaseOverlapExpansionSet):
221 def __init__(self, msoe, local_indices, Mmystart, mynao):
222 self.msoe = msoe
223 self.local_indices = local_indices
224 BaseOverlapExpansionSet.__init__(self, msoe.shape)
226 self.Mmystart = Mmystart
227 self.mynao = mynao
229 M_a = msoe.M1_a
230 natoms = len(M_a)
231 a = 0
232 while a < natoms and M_a[a] <= Mmystart:
233 a += 1
234 a -= 1
235 self.astart = a
237 while a < natoms and M_a[a] < Mmystart + mynao:
238 a += 1
239 self.aend = a
241 def evaluate_slice(self, disp, x_xqNM):
242 a1 = disp.a1
243 a2 = disp.a2
244 if a2 in self.local_indices and (self.astart <= a1 < self.aend):
245 # assert a2 <= a1
246 msoe = self.msoe
247 I1 = msoe.I1_a[a1]
248 I2 = msoe.I2_a[a2]
249 tsoe = msoe.tsoe_II[I1, I2]
250 x_qxmm = tsoe.zeros(x_xqNM.shape[:-2], dtype=x_xqNM.dtype)
251 disp.evaluate_overlap(tsoe, x_qxmm)
252 Mstart1 = msoe.M1_a[a1] - self.Mmystart
253 Mend1 = Mstart1 + tsoe.shape[0]
254 Mstart1b = max(0, Mstart1)
255 Mend1b = min(self.mynao, Mend1)
256 Mstart2 = msoe.M2_a[a2]
257 Mend2 = Mstart2 + tsoe.shape[1]
258 x_xqNM[..., Mstart1b:Mend1b, Mstart2:Mend2] += \
259 x_qxmm[..., Mstart1b - Mstart1:Mend1b - Mstart1, :]
260 # This is all right as long as we are only interested in a2 <= a1
261 # if a1 in self.local_indices and a2 < a1 and (self.astart <=
262 # a2 < self.aend):
263 # self.evaluate_slice(disp.reverse(), x_xqNM)
266class NeighborPairs:
267 """Class for looping over pairs of atoms using a neighbor list."""
268 def __init__(self, cutoff_a, cell_cv, pbc_c, self_interaction):
269 self.neighbors = PrimitiveNeighborList(
270 cutoff_a, skin=0, sorted=True,
271 self_interaction=self_interaction,
272 use_scaled_positions=True)
273 self.cell_cv = cell_cv
274 self.pbc_c = pbc_c
276 def set_positions(self, spos_ac):
277 self.spos_ac = spos_ac
278 self.neighbors.update(self.pbc_c, self.cell_cv, spos_ac)
280 def iter(self):
281 cell_cv = self.cell_cv
282 for a1, spos1_c in enumerate(self.spos_ac):
283 a2_a, offsets = self.neighbors.get_neighbors(a1)
284 for a2, offset in zip(a2_a, offsets):
285 spos2_c = self.spos_ac[a2] + offset
286 R_c = np.dot(spos2_c - spos1_c, cell_cv)
287 yield a1, a2, R_c, offset
290class PairFilter:
291 def __init__(self, pairs):
292 self.pairs = pairs
294 def set_positions(self, spos_ac):
295 self.pairs.set_positions(spos_ac)
297 def iter(self):
298 return self.pairs.iter()
301class PairsWithSelfinteraction(PairFilter):
302 def iter(self):
303 for a1, a2, R_c, offset in self.pairs.iter():
304 yield a1, a2, R_c, offset
305 if a1 == a2 and offset.any():
306 yield a1, a1, -R_c, -offset
309class OppositeDirection(PairFilter):
310 def iter(self):
311 for a1, a2, R_c, offset in self.pairs.iter():
312 yield a2, a1, -R_c, -offset
315class FourierTransformer(FourierBesselTransformer):
316 def calculate_overlap_expansion(self, phit1phit2_q, l1, l2):
317 """Calculate list of splines for one overlap integral.
319 Given two Fourier transformed functions, return list of splines
320 in real space necessary to evaluate their overlap.
322 phi (q) * phi (q) --> [phi (r), ..., phi (r)] .
323 l1 l2 lmin lmax
325 The overlap <phi1 | phi2> can then be calculated by linear
326 combinations of the returned splines with appropriate Gaunt
327 coefficients.
328 """
329 lmax = l1 + l2
330 splines = []
331 R = np.arange(self.Q) * self.dr
332 R1 = R.copy()
333 R1[0] = 1.0
334 k1 = self.k_q.copy()
335 k1[0] = 1.0
336 a_q = phit1phit2_q
337 for l in range(lmax % 2, lmax + 1, 2):
338 a_g = (8 * ffbt(l, a_q * k1**(-2 - lmax - l), self.k_q, R) /
339 R1**(2 * l + 1))
340 if l == 0:
341 a_g[0] = 8 * np.sum(a_q * k1**(-lmax)) * self.dk
342 else:
343 a_g[0] = a_g[1] # XXXX
344 a_g *= (-1)**((l1 - l2 - l) // 2)
345 n = len(a_g) // 256
346 s = Spline.from_data(
347 l, 2 * self.rcut, np.concatenate((a_g[::n], [0.0])),
348 )
349 splines.append(s)
350 return OverlapExpansion(l1, l2, splines)
352 def laplacian(self, f_jq):
353 return 0.5 * f_jq * self.k_q**2.0
356class TwoSiteOverlapCalculator:
357 def __init__(self, transformer):
358 self.transformer = transformer
360 def transform(self, f_j):
361 f_jq = np.array([self.transformer.transform(f) for f in f_j])
362 return f_jq
364 def calculate_expansions(self, la_j, fa_jq, lb_j, fb_jq):
365 nja = len(la_j)
366 njb = len(lb_j)
367 assert nja == len(fa_jq) and njb == len(fb_jq)
368 oe_jj = np.empty((nja, njb), dtype=object)
369 for ja, (la, fa_q) in enumerate(zip(la_j, fa_jq)):
370 for jb, (lb, fb_q) in enumerate(zip(lb_j, fb_jq)):
371 a_q = fa_q * fb_q
372 oe = self.transformer.calculate_overlap_expansion(a_q, la, lb)
373 oe_jj[ja, jb] = oe
374 return TwoSiteOverlapExpansions(la_j, lb_j, oe_jj)
376 def calculate_kinetic_expansions(self, l_j, f_jq):
377 t_jq = self.transformer.laplacian(f_jq)
378 return self.calculate_expansions(l_j, f_jq, l_j, t_jq)
380 def laplacian(self, f_jq):
381 t_jq = self.transformer.laplacian(f_jq)
382 return t_jq
385class ManySiteOverlapCalculator:
386 def __init__(self, twosite_calculator, I1_a, I2_a):
387 """Create VeryManyOverlaps object.
389 twosite_calculator: instance of TwoSiteOverlapCalculator
390 I_a: mapping from atom index (as in spos_a) to unique atom type"""
391 self.twosite_calculator = twosite_calculator
392 self.I1_a = I1_a
393 self.I2_a = I2_a
395 def transform(self, f_Ij):
396 f_Ijq = [self.twosite_calculator.transform(f_j) for f_j in f_Ij]
397 return f_Ijq
399 def calculate_expansions(self, l1_Ij, f1_Ijq, l2_Ij, f2_Ijq):
400 # Naive implementation, just loop over everything
401 # We should only need half of them
402 nI1 = len(l1_Ij)
403 nI2 = len(l2_Ij)
404 assert len(l1_Ij) == len(f1_Ijq) and len(l2_Ij) == len(f2_Ijq)
405 tsoe_II = np.empty((nI1, nI2), dtype=object)
406 calc = self.twosite_calculator
407 for I1, (l1_j, f1_jq) in enumerate(zip(l1_Ij, f1_Ijq)):
408 for I2, (l2_j, f2_jq) in enumerate(zip(l2_Ij, f2_Ijq)):
409 tsoe = calc.calculate_expansions(l1_j, f1_jq, l2_j, f2_jq)
410 tsoe_II[I1, I2] = tsoe
411 return ManySiteOverlapExpansions(tsoe_II, self.I1_a, self.I2_a)
413 def calculate_kinetic_expansions(self, l_Ij, f_Ijq):
414 t_Ijq = [self.twosite_calculator.laplacian(f_jq) for f_jq in f_Ijq]
415 return self.calculate_expansions(l_Ij, f_Ijq, l_Ij, t_Ijq)
418class AtomicDisplacement:
419 def __init__(self, factory, a1, a2, R_c, offset, phases):
420 self.factory = factory
421 self.a1 = a1
422 self.a2 = a2
423 self.R_c = R_c
424 self.offset = offset
425 self.phases = phases
426 self.r = np.linalg.norm(R_c)
427 self._set_spherical_harmonics(R_c)
429 def _set_spherical_harmonics(self, R_c):
430 self.rlY_lm = LazySphericalHarmonics(R_c)
432 # XXX new
433 def evaluate_direct(self, oe, dst_xqmm):
434 src_xmm = self.evaluate_direct_without_phases(oe)
435 self.phases.apply(src_xmm, dst_xqmm)
437 # XXX new
438 def evaluate_direct_without_phases(self, oe):
439 return oe.evaluate(self.r, self.rlY_lm)
441 # XXX clean up unnecessary methods
442 def overlap_without_phases(self, oe):
443 return oe.evaluate(self.r, self.rlY_lm)
445 def _evaluate_without_phases(self, oe):
446 return self.overlap_without_phases(oe)
448 def evaluate_overlap(self, oe, dst_xqmm):
449 src_xmm = self._evaluate_without_phases(oe)
450 timer.start('phases')
451 self.phases.apply(src_xmm, dst_xqmm)
452 timer.stop('phases')
454 def reverse(self):
455 return self.factory.displacementclass(self.factory, self.a2, self.a1,
456 -self.R_c, -self.offset,
457 self.phases.inverse())
460class LazySphericalHarmonics:
461 """Class for caching spherical harmonics as they are calculated.
463 Behaves like a list Y_lm, but really calculates (or retrieves) Y_m
464 once a given value of l is __getitem__'d."""
465 def __init__(self, R_c):
466 self.R_c = np.asarray(R_c)
467 self.Y_lm = {}
468 self.Y_lm1 = np.empty(0)
469 # self.dYdr_lmc = {}
470 self.lmax = 0
472 def evaluate(self, l):
473 rlY_m = np.empty(2 * l + 1)
474 Yl(l, self.R_c, rlY_m)
475 return rlY_m
477 def __getitem__(self, l):
478 Y_m = self.Y_lm.get(l)
479 if Y_m is None:
480 Y_m = self.evaluate(l)
481 self.Y_lm[l] = Y_m
482 return Y_m
484 def toarray(self, lmax):
485 if lmax > self.lmax:
486 self.Y_lm1 = np.concatenate([self.Y_lm1] +
487 [self.evaluate(l).ravel()
488 for l in range(self.lmax, lmax)])
489 self.lmax = lmax
490 return self.Y_lm1
493class LazySphericalHarmonicsDerivative(LazySphericalHarmonics):
494 def evaluate(self, l):
495 drlYdR_mc = np.empty((2 * l + 1, 3))
496 L0 = l**2
497 for m in range(2 * l + 1):
498 drlYdR_mc[m, :] = nablarlYL(L0 + m, self.R_c)
499 return drlYdR_mc
502class DerivativeAtomicDisplacement(AtomicDisplacement):
503 def _set_spherical_harmonics(self, R_c):
504 self.rlY_lm = LazySphericalHarmonics(R_c)
505 self.drlYdr_lmc = LazySphericalHarmonicsDerivative(R_c)
507 if R_c.any():
508 self.Rhat_c = R_c / self.r
509 else:
510 self.Rhat_c = np.zeros(3)
512 def derivative_without_phases(self, oe):
513 return oe.derivative(self.r, self.Rhat_c, self.rlY_lm, self.drlYdr_lmc)
515 def _evaluate_without_phases(self, oe): # override
516 return self.derivative_without_phases(oe)
519class NullPhases:
520 def __init__(self, ibzk_qc, offset):
521 pass
523 def apply(self, src_xMM, dst_qxMM):
524 assert len(dst_qxMM) == 1
525 dst_qxMM[0][:] += src_xMM
527 def inverse(self):
528 return self
531class BlochPhases:
532 def __init__(self, ibzk_qc, offset):
533 self.phase_q = np.exp(-2j * pi * np.dot(ibzk_qc, offset))
534 self.ibzk_qc = ibzk_qc
535 self.offset = offset
537 def apply(self, src_xMM, dst_qxMM):
538 assert dst_qxMM.dtype == complex, dst_qxMM.dtype
539 for phase, dst_xMM in zip(self.phase_q, dst_qxMM):
540 dst_xMM[:] += phase * src_xMM
542 def inverse(self):
543 return BlochPhases(-self.ibzk_qc, self.offset)
546class TwoCenterIntegralCalculator:
547 # This class knows how to apply phases, and whether to call the
548 # various derivative() or evaluate() methods
549 def __init__(self, ibzk_qc=None, derivative=False):
550 if derivative:
551 displacementclass = DerivativeAtomicDisplacement
552 else:
553 displacementclass = AtomicDisplacement
554 self.displacementclass = displacementclass
556 if ibzk_qc is None or not ibzk_qc.any():
557 self.phaseclass = NullPhases
558 else:
559 self.phaseclass = BlochPhases
560 self.ibzk_qc = ibzk_qc
561 self.derivative = derivative
563 def calculate(self, atompairs, expansions, arrays):
564 for disp in self.iter(atompairs):
565 for expansion, X_qxMM in zip(expansions, arrays):
566 expansion.evaluate_slice(disp, X_qxMM)
568 def iter(self, atompairs):
569 for a1, a2, R_c, offset in atompairs.iter():
570 # if a1 == a2 and self.derivative:
571 # continue
572 phase_applier = self.phaseclass(self.ibzk_qc, offset)
573 yield self.displacementclass(self, a1, a2, R_c, offset,
574 phase_applier)
577class NewTwoCenterIntegrals:
578 def __init__(self, cell_cv, pbc_c, setups, ibzk_qc, gamma):
579 self.cell_cv = cell_cv
580 self.pbc_c = pbc_c
581 self.ibzk_qc = ibzk_qc
582 self.gamma = gamma
584 timer.start('tci init')
585 cutoff_I = []
586 setups_I = setups.setups.values()
587 I_setup = {}
588 for I, setup in enumerate(setups_I):
589 I_setup[setup] = I
590 cutoff_I.append(max([func.get_cutoff()
591 for func
592 in setup.basis_functions_J + setup.pt_j]))
594 I_a = []
595 for setup in setups:
596 I_a.append(I_setup[setup])
598 cutoff_a = [cutoff_I[I] for I in I_a]
600 self.cutoff_a = cutoff_a # convenient for writing the new new overlap
601 self.I_a = I_a
602 self.setups_I = setups_I
603 self.atompairs = PairsWithSelfinteraction(NeighborPairs(cutoff_a,
604 cell_cv,
605 pbc_c,
606 True))
608 scale = 0.01 # XXX minimal distance scale
609 cutoff_close_a = [covalent_radii[int(s.Z)] / Bohr * scale
610 for s in setups]
611 self.atoms_close = NeighborPairs(cutoff_close_a, cell_cv, pbc_c, False)
613 rcut = max(cutoff_I + [0.001])
614 transformer = FourierTransformer(rcut, N=2**10)
615 tsoc = TwoSiteOverlapCalculator(transformer)
616 self.msoc = ManySiteOverlapCalculator(tsoc, I_a, I_a)
617 self.calculate_expansions()
619 self.calculate = self.evaluate # XXX compatibility
621 self.set_matrix_distribution(None, None)
622 timer.stop('tci init')
624 def set_matrix_distribution(self, Mmystart, mynao):
625 """Distribute matrices using BLACS."""
626 # Range of basis functions for BLACS distribution of matrices:
627 self.Mmystart = Mmystart
628 self.mynao = mynao
629 self.blacs = mynao is not None
631 def calculate_expansions(self):
632 timer.start('tci calc exp')
633 phit_Ij = [setup.basis_functions_J for setup in self.setups_I]
634 l_Ij = []
635 for phit_j in phit_Ij:
636 l_Ij.append([phit.get_angular_momentum_number()
637 for phit in phit_j])
639 pt_l_Ij = [setup.l_j for setup in self.setups_I]
640 pt_Ij = [setup.pt_j for setup in self.setups_I]
641 phit_Ijq = self.msoc.transform(phit_Ij)
642 pt_Ijq = self.msoc.transform(pt_Ij)
644 msoc = self.msoc
646 self.Theta_expansions = msoc.calculate_expansions(l_Ij, phit_Ijq,
647 l_Ij, phit_Ijq)
648 self.T_expansions = msoc.calculate_kinetic_expansions(l_Ij, phit_Ijq)
649 self.P_expansions = msoc.calculate_expansions(l_Ij, phit_Ijq,
650 pt_l_Ij, pt_Ijq)
651 timer.stop('tci calc exp')
653 def _calculate(self, calc, spos_ac, Theta_qxMM, T_qxMM, P_aqxMi):
654 Theta_qxMM.fill(0.0)
655 T_qxMM.fill(0.0)
656 for P_qxMi in P_aqxMi.values():
657 P_qxMi.fill(0.0)
659 if 1: # XXX
660 self.atoms_close.set_positions(spos_ac)
661 wrk = [x for x in self.atoms_close.iter()]
662 if len(wrk) != 0:
663 txt = ''
664 for a1, a2, R_c, offset in wrk:
665 txt += 'Atom %d and Atom %d in cell (%d, %d, %d)\n' % \
666 (a1, a2, offset[0], offset[1], offset[2])
667 raise RuntimeError('Atoms too close!\n' + txt)
669 self.atompairs.set_positions(spos_ac)
671 if self.blacs:
672 # S and T matrices are distributed:
673 expansions = [
674 BlacsOverlapExpansions(self.Theta_expansions,
675 P_aqxMi, self.Mmystart, self.mynao),
676 BlacsOverlapExpansions(self.T_expansions,
677 P_aqxMi, self.Mmystart, self.mynao)]
678 else:
679 expansions = [DomainDecomposedExpansions(self.Theta_expansions,
680 P_aqxMi),
681 DomainDecomposedExpansions(self.T_expansions,
682 P_aqxMi)]
684 expansions.append(ManySiteDictionaryWrapper(self.P_expansions,
685 P_aqxMi))
686 arrays = [Theta_qxMM, T_qxMM, P_aqxMi]
687 timer.start('tci calculate')
688 calc.calculate(OppositeDirection(self.atompairs), expansions, arrays)
689 timer.stop('tci calculate')
691 def evaluate(self, spos_ac, Theta_qMM, T_qMM, P_aqMi):
692 calc = TwoCenterIntegralCalculator(self.ibzk_qc, derivative=False)
693 self._calculate(calc, spos_ac, Theta_qMM, T_qMM, P_aqMi)
694 if not self.blacs:
695 for X_MM in list(Theta_qMM) + list(T_qMM):
696 tri2full(X_MM, UL=UL)
698 def derivative(self, spos_ac, dThetadR_qcMM, dTdR_qcMM, dPdR_aqcMi):
699 calc = TwoCenterIntegralCalculator(self.ibzk_qc, derivative=True)
700 self._calculate(calc, spos_ac, dThetadR_qcMM, dTdR_qcMM, dPdR_aqcMi)
702 def antihermitian(src, dst):
703 np.conj(-src, dst)
705 if not self.blacs:
706 for X_cMM in list(dThetadR_qcMM) + list(dTdR_qcMM):
707 for X_MM in X_cMM:
708 tri2full(X_MM, UL=UL, map=antihermitian)
710 calculate_derivative = derivative # XXX compatibility
712 def estimate_memory(self, mem):
713 mem.subnode('TCI not impl.', 0)