1import sys 2from time import ctime 3from pathlib import Path 4import pickle 5 6import numpy as np 7 8from ase.units import Hartree 9from gpaw.utilities import convert_string_to_fd 10from ase.utils.timing import Timer, timer 11 12import gpaw.mpi as mpi 13from gpaw.blacs import BlacsGrid, BlacsDescriptor, Redistributor 14from gpaw.response.kspair import get_calc 15from gpaw.response.kslrf import FrequencyDescriptor 16from gpaw.response.chiks import ChiKS 17from gpaw.response.kxc import get_fxc 18from gpaw.response.kernels import get_coulomb_kernel 19 20 21class FourComponentSusceptibilityTensor: 22 """Class calculating the full four-component susceptibility tensor""" 23 24 def __init__(self, gs, fxc='ALDA', fxckwargs={}, 25 eta=0.2, ecut=50, gammacentered=False, 26 disable_point_group=True, disable_time_reversal=True, 27 bandsummation='pairwise', nbands=None, 28 bundle_integrals=True, bundle_kptpairs=False, 29 world=mpi.world, nblocks=1, txt=sys.stdout): 30 """ 31 Currently, everything is in plane wave mode. 32 If additional modes are implemented, maybe look to fxc to see how 33 multiple modes can be supported. 34 35 Parameters 36 ---------- 37 gs : see gpaw.response.chiks, gpaw.response.kslrf 38 fxc, fxckwargs : see gpaw.response.fxc 39 eta, ecut, gammacentered 40 disable_point_group, 41 disable_time_reversal, 42 bandsummation, nbands, 43 bundle_integrals, bundle_kptpairs, 44 world, nblocks, txt : see gpaw.response.chiks, gpaw.response.kslrf 45 """ 46 # Initiate output file and timer 47 self.world = world 48 self.fd = convert_string_to_fd(txt, world) 49 self.cfd = self.fd 50 self.timer = Timer() 51 52 # Load ground state calculation 53 self.calc = get_calc(gs, fd=self.fd, timer=self.timer) 54 55 # The plane wave basis is defined by keywords 56 self.ecut = None if ecut is None else ecut / Hartree 57 self.gammacentered = gammacentered 58 59 # Initiate Kohn-Sham susceptibility and fxc objects 60 self.chiks = ChiKS(self.calc, eta=eta, ecut=ecut, 61 gammacentered=gammacentered, 62 disable_point_group=disable_point_group, 63 disable_time_reversal=disable_time_reversal, 64 bandsummation=bandsummation, nbands=nbands, 65 bundle_integrals=bundle_integrals, 66 bundle_kptpairs=bundle_kptpairs, world=world, 67 nblocks=nblocks, txt=self.fd, timer=self.timer) 68 self.fxc = get_fxc(self.calc, fxc, 69 response='susceptibility', mode='pw', 70 world=self.chiks.world, txt=self.chiks.fd, 71 timer=self.timer, **fxckwargs) 72 73 # Parallelization over frequencies depends on the frequency input 74 self.mynw = None 75 self.w1 = None 76 self.w2 = None 77 78 def get_macroscopic_component(self, spincomponent, q_c, frequencies, 79 filename=None, txt=None): 80 """Calculates the spatially averaged (macroscopic) component of the 81 susceptibility tensor and writes it to a file. 82 83 Parameters 84 ---------- 85 spincomponent, q_c, 86 frequencies : see gpaw.response.chiks, gpaw.response.kslrf 87 filename : str 88 Save chiks_w and chi_w to file of given name. 89 Defaults to: 90 'chi%s_q«%+d-%+d-%+d».csv' % (spincomponent, 91 *tuple((q_c * kd.N_c).round())) 92 txt : str 93 Save output of the calculation of this specific component into 94 a file with the filename of the given input. 95 96 Returns 97 ------- 98 see calculate_macroscopic_component 99 """ 100 101 if filename is None: 102 tup = (spincomponent, 103 *tuple((q_c * self.calc.wfs.kd.N_c).round())) 104 filename = 'chi%s_q«%+d-%+d-%+d».csv' % tup 105 106 (omega_w, 107 chiks_w, 108 chi_w) = self.calculate_macroscopic_component(spincomponent, q_c, 109 frequencies, 110 txt=txt) 111 112 write_macroscopic_component(omega_w, chiks_w, chi_w, 113 filename, self.world) 114 115 return omega_w, chiks_w, chi_w 116 117 def calculate_macroscopic_component(self, spincomponent, 118 q_c, frequencies, txt=None): 119 """Calculates the spatially averaged (macroscopic) component of the 120 susceptibility tensor. 121 122 Parameters 123 ---------- 124 spincomponent, q_c, 125 frequencies : see gpaw.response.chiks, gpaw.response.kslrf 126 txt : see get_macroscopic_component 127 128 Returns 129 ------- 130 omega_w, chiks_w, chi_w : nd.array, nd.array, nd.array 131 omega_w: frequencies in eV 132 chiks_w: macroscopic dynamic susceptibility (Kohn-Sham system) 133 chi_w: macroscopic dynamic susceptibility 134 """ 135 (pd, wd, 136 chiks_wGG, chi_wGG) = self.calculate_component(spincomponent, q_c, 137 frequencies, txt=txt) 138 139 # Macroscopic component 140 chiks_w = chiks_wGG[:, 0, 0] 141 chi_w = chi_wGG[:, 0, 0] 142 143 # Collect data for all frequencies 144 omega_w = wd.get_data() * Hartree 145 chiks_w = self.collect(chiks_w) 146 chi_w = self.collect(chi_w) 147 148 return omega_w, chiks_w, chi_w 149 150 def get_component_array(self, spincomponent, q_c, frequencies, 151 array_ecut=50, filename=None, txt=None): 152 """Calculates a specific spin component of the susceptibility tensor, 153 collects it as a numpy array in a reduced plane wave description 154 and writes it to a file. 155 156 Parameters 157 ---------- 158 spincomponent, q_c, 159 frequencies : see gpaw.response.chiks, gpaw.response.kslrf 160 array_ecut : see calculate_component_array 161 filename : str 162 Save chiks_w and chi_w to pickle file of given name. 163 Defaults to: 164 'chi%sGG_q«%+d-%+d-%+d».pckl' % (spincomponent, 165 *tuple((q_c * kd.N_c).round())) 166 txt : str 167 Save output of the calculation of this specific component into 168 a file with the filename of the given input. 169 170 Returns 171 ------- 172 see calculate_component_array 173 """ 174 175 if filename is None: 176 tup = (spincomponent, 177 *tuple((q_c * self.calc.wfs.kd.N_c).round())) 178 filename = 'chi%sGG_q«%+d-%+d-%+d».pckl' % tup 179 180 (omega_w, G_Gc, chiks_wGG, 181 chi_wGG) = self.calculate_component_array(spincomponent, 182 q_c, 183 frequencies, 184 array_ecut=array_ecut, 185 txt=txt) 186 187 # Write susceptibilities to a pickle file 188 write_component(omega_w, G_Gc, chiks_wGG, chi_wGG, 189 filename, self.world) 190 191 return omega_w, G_Gc, chiks_wGG, chi_wGG 192 193 def calculate_component_array(self, spincomponent, q_c, frequencies, 194 array_ecut=50, txt=None): 195 """Calculates a specific spin component of the susceptibility tensor 196 and collects it as a numpy array in a reduced plane wave description. 197 198 Parameters 199 ---------- 200 spincomponent, q_c, 201 frequencies : see gpaw.response.chiks, gpaw.response.kslrf 202 array_ecut : float 203 Energy cutoff for the reduced plane wave representation. 204 The susceptibility is returned in the reduced representation. 205 206 Returns 207 ------- 208 omega_w, G_Gc, chiks_wGG, chi_wGG : nd.array(s) 209 omega_w: frequencies in eV 210 G_Gc : plane wave repr. as coordinates on the reciprocal lattice 211 chiks_wGG: dynamic susceptibility (Kohn-Sham system) 212 chi_wGG: dynamic susceptibility 213 """ 214 (pd, wd, 215 chiks_wGG, chi_wGG) = self.calculate_component(spincomponent, q_c, 216 frequencies, txt=txt) 217 218 # Get frequencies in eV 219 omega_w = wd.get_data() * Hartree 220 221 # Get susceptibility in a reduced plane wave representation 222 mask_G = get_pw_reduction_map(pd, array_ecut) 223 chiks_wGG = np.ascontiguousarray(chiks_wGG[:, mask_G, :][:, :, mask_G]) 224 chi_wGG = np.ascontiguousarray(chi_wGG[:, mask_G, :][:, :, mask_G]) 225 226 # Get reduced plane wave repr. as coordinates on the reciprocal lattice 227 G_Gc = get_pw_coordinates(pd)[mask_G] 228 229 # Gather susceptibilities for all frequencies 230 chiks_wGG = self.gather(chiks_wGG, wd) 231 chi_wGG = self.gather(chi_wGG, wd) 232 233 return omega_w, G_Gc, chiks_wGG, chi_wGG 234 235 def calculate_component(self, spincomponent, q_c, frequencies, txt=None): 236 """Calculate a single component of the susceptibility tensor. 237 238 Parameters 239 ---------- 240 spincomponent, q_c, 241 frequencies : see gpaw.response.chiks, gpaw.response.kslrf 242 243 Returns 244 ------- 245 pd : PWDescriptor 246 Descriptor object for the plane wave basis 247 wd : FrequencyDescriptor 248 Descriptor object for the calculated frequencies 249 chiks_wGG : ndarray 250 The process' block of the Kohn-Sham susceptibility component 251 chi_wGG : ndarray 252 The process' block of the full susceptibility component 253 """ 254 pd = self.get_PWDescriptor(q_c) 255 256 # Initiate new call-output file, if supplied 257 if txt is not None: 258 # Store timer and close old call-output file 259 self.write_timer() 260 if str(self.fd) != str(self.cfd): 261 print('\nClosing, %s' % ctime(), file=self.cfd) 262 self.cfd.close() 263 # Initiate new output file 264 self.cfd = convert_string_to_fd(txt, self.world) 265 # Print to output file 266 if str(self.fd) != str(self.cfd) or txt is not None: 267 print('---------------', file=self.fd) 268 print(f'Calculating susceptibility spincomponent={spincomponent}' 269 f'with q_c={pd.kd.bzk_kc[0]}', file=self.fd) 270 271 wd = FrequencyDescriptor(np.asarray(frequencies) / Hartree) 272 273 # Initialize parallelization over frequencies 274 nw = len(wd) 275 self.mynw = (nw + self.world.size - 1) // self.world.size 276 self.w1 = min(self.mynw * self.world.rank, nw) 277 self.w2 = min(self.w1 + self.mynw, nw) 278 279 return self._calculate_component(spincomponent, pd, wd) 280 281 def get_PWDescriptor(self, q_c): 282 """Get the planewave descriptor defining the plane wave basis for the 283 given momentum transfer q_c.""" 284 from gpaw.kpt_descriptor import KPointDescriptor 285 from gpaw.wavefunctions.pw import PWDescriptor 286 q_c = np.asarray(q_c, dtype=float) 287 qd = KPointDescriptor([q_c]) 288 pd = PWDescriptor(self.ecut, self.calc.wfs.gd, 289 complex, qd, gammacentered=self.gammacentered) 290 return pd 291 292 def _calculate_component(self, spincomponent, pd, wd): 293 """In-place calculation of the given spin-component.""" 294 if spincomponent in ['+-', '-+']: 295 # No Hartree kernel 296 K_GG = self.fxc(spincomponent, pd, txt=self.cfd) 297 else: 298 # Calculate Hartree kernel 299 Kbare_G = get_coulomb_kernel(pd, self.calc.wfs.kd.N_c) 300 vsqrt_G = Kbare_G ** 0.5 301 K_GG = np.eye(len(vsqrt_G)) * vsqrt_G * vsqrt_G[:, np.newaxis] 302 303 # Calculate exchange-correlation kernel 304 Kxc_GG = self.fxc(spincomponent, pd, txt=self.cfd) 305 if Kxc_GG is not None: 306 K_GG += Kxc_GG 307 308 chiks_wGG = self.calculate_ks_component(spincomponent, pd, 309 wd, txt=self.cfd) 310 311 chi_wGG = self.invert_dyson(chiks_wGG, K_GG) 312 313 return pd, wd, chiks_wGG, chi_wGG 314 315 def write_timer(self): 316 """Write timer to call-output file and initiate a new.""" 317 self.timer.write(self.cfd) 318 self.timer = Timer() 319 320 # Update all other class instance timers 321 self.chiks.timer = self.timer 322 self.chiks.integrator.timer = self.timer 323 self.chiks.kspair.timer = self.timer 324 self.chiks.pme.timer = self.timer 325 self.fxc.timer = self.timer 326 327 def calculate_ks_component(self, spincomponent, pd, wd, txt=None): 328 """Calculate a single component of the Kohn-Sham susceptibility tensor. 329 330 Parameters 331 ---------- 332 spincomponent : see gpaw.response.chiks, gpaw.response.kslrf 333 pd, wd : see calculate_component 334 335 Returns 336 ------- 337 chiks_wGG : ndarray 338 The process' block of the Kohn-Sham susceptibility component 339 """ 340 # ChiKS calculates the susceptibility distributed over plane waves 341 _, chiks_wGG = self.chiks.calculate(pd, wd, txt=txt, 342 spincomponent=spincomponent) 343 344 # Redistribute memory, so each block has its own frequencies, but all 345 # plane waves (for easy invertion of the Dyson-like equation) 346 chiks_wGG = self.distribute_frequencies(chiks_wGG) 347 348 return chiks_wGG 349 350 @timer('Invert dyson-like equation') 351 def invert_dyson(self, chiks_wGG, Khxc_GG): 352 """Invert the Dyson-like equation: 353 354 chi = chi_ks + chi_ks Khxc chi 355 """ 356 chi_wGG = np.empty_like(chiks_wGG) 357 for w, chiks_GG in enumerate(chiks_wGG): 358 chi_GG = np.dot(np.linalg.inv(np.eye(len(chiks_GG)) + 359 np.dot(chiks_GG, Khxc_GG)), 360 chiks_GG) 361 362 chi_wGG[w] = chi_GG 363 364 return chi_wGG 365 366 def collect(self, a_w): 367 """Collect frequencies from all blocks""" 368 # More documentation is needed! XXX 369 world = self.chiks.world 370 b_w = np.zeros(self.mynw, a_w.dtype) 371 b_w[:self.w2 - self.w1] = a_w 372 nw = len(self.chiks.omega_w) 373 A_w = np.empty(world.size * self.mynw, a_w.dtype) 374 world.all_gather(b_w, A_w) 375 return A_w[:nw] 376 377 def gather(self, A_wGG, wd): 378 """Gather a full susceptibility array to root.""" 379 # Allocate arrays to gather (all need to be the same shape) 380 shape = (self.mynw,) + A_wGG.shape[1:] 381 tmp_wGG = np.empty(shape, dtype=A_wGG.dtype) 382 tmp_wGG[:self.w2 - self.w1] = A_wGG 383 384 # Allocate array for the gathered data 385 if self.world.rank == 0: 386 # Make room for all frequencies 387 shape = (self.mynw * self.world.size,) + A_wGG.shape[1:] 388 allA_wGG = np.empty(shape, dtype=A_wGG.dtype) 389 else: 390 allA_wGG = None 391 392 self.world.gather(tmp_wGG, 0, allA_wGG) 393 394 # Return array for w indeces on frequency grid 395 if allA_wGG is not None: 396 allA_wGG = allA_wGG[:len(wd), :, :] 397 398 return allA_wGG 399 400 @timer('Distribute frequencies') 401 def distribute_frequencies(self, chiks_wGG): 402 """Distribute frequencies to all cores.""" 403 # More documentation is needed! XXX 404 world = self.chiks.world 405 comm = self.chiks.interblockcomm 406 407 if world.size == 1: 408 return chiks_wGG 409 410 nw = len(self.chiks.omega_w) 411 nG = chiks_wGG.shape[2] 412 mynw = (nw + world.size - 1) // world.size 413 mynG = (nG + comm.size - 1) // comm.size 414 415 wa = min(world.rank * mynw, nw) 416 wb = min(wa + mynw, nw) 417 418 if self.chiks.interblockcomm.size == 1: 419 return chiks_wGG[wa:wb].copy() 420 421 if self.chiks.intrablockcomm.rank == 0: 422 bg1 = BlacsGrid(comm, 1, comm.size) 423 in_wGG = chiks_wGG.reshape((nw, -1)) 424 else: 425 bg1 = BlacsGrid(None, 1, 1) 426 # bg1 = DryRunBlacsGrid(mpi.serial_comm, 1, 1) 427 in_wGG = np.zeros((0, 0), complex) 428 md1 = BlacsDescriptor(bg1, nw, nG**2, nw, mynG * nG) 429 430 bg2 = BlacsGrid(world, world.size, 1) 431 md2 = BlacsDescriptor(bg2, nw, nG**2, mynw, nG**2) 432 433 r = Redistributor(world, md1, md2) 434 shape = (wb - wa, nG, nG) 435 out_wGG = np.empty(shape, complex) 436 r.redistribute(in_wGG, out_wGG.reshape((wb - wa, nG**2))) 437 438 return out_wGG 439 440 def close(self): 441 self.timer.write(self.cfd) 442 print('\nClosing, %s' % ctime(), file=self.cfd) 443 self.cfd.close() 444 print('\nClosing, %s' % ctime(), file=self.fd) 445 self.fd.close() 446 447 448def get_pw_reduction_map(pd, ecut): 449 """Get a mask to reduce the plane wave representation. 450 451 Please remark, that the response code currently works with one q-vector 452 at a time, at thus only a single plane wave representation at a time. 453 454 Returns 455 ------- 456 mask_G : nd.array (dtype=bool) 457 Mask which reduces the representation 458 """ 459 assert ecut is not None 460 ecut /= Hartree 461 assert ecut <= pd.ecut 462 463 # List of all plane waves 464 G_Gv = np.array([pd.G_Qv[Q] for Q in pd.Q_qG[0]]) 465 466 if pd.gammacentered: 467 mask_G = ((G_Gv ** 2).sum(axis=1) <= 2 * ecut) 468 else: 469 mask_G = (((G_Gv + pd.K_qv[0]) ** 2).sum(axis=1) <= 2 * ecut) 470 471 return mask_G 472 473 474def get_pw_coordinates(pd): 475 """Get the reciprocal lattice vector coordinates corresponding to a 476 givne plane wave basis. 477 478 Please remark, that the response code currently works with one q-vector 479 at a time, at thus only a single plane wave representation at a time. 480 481 Returns 482 ------- 483 G_Gc : nd.array (dtype=int) 484 Coordinates on the reciprocal lattice 485 """ 486 # List of all plane waves 487 G_Gv = np.array([pd.G_Qv[Q] for Q in pd.Q_qG[0]]) 488 489 # Use cell to get coordinates 490 B_cv = 2.0 * np.pi * pd.gd.icell_cv 491 return np.round(np.dot(G_Gv, np.linalg.inv(B_cv))).astype(int) 492 493 494def write_macroscopic_component(omega_w, chiks_w, chi_w, filename, world): 495 """Write the spatially averaged dynamic susceptibility.""" 496 assert isinstance(filename, str) 497 if world.rank == 0: 498 with Path(filename).open('w') as fd: 499 for omega, chiks, chi in zip(omega_w, chiks_w, chi_w): 500 print('%.6f, %.6f, %.6f, %.6f, %.6f' % 501 (omega, chiks.real, chiks.imag, chi.real, chi.imag), 502 file=fd) 503 504 505def read_macroscopic_component(filename): 506 """Read a stored macroscopic susceptibility file""" 507 d = np.loadtxt(filename, delimiter=', ') 508 omega_w = d[:, 0] 509 chiks_w = np.array(d[:, 1], complex) 510 chiks_w.imag = d[:, 2] 511 chi_w = np.array(d[:, 3], complex) 512 chi_w.imag = d[:, 4] 513 514 return omega_w, chiks_w, chi_w 515 516 517def write_component(omega_w, G_Gc, chiks_wGG, chi_wGG, filename, world): 518 """Write the dynamic susceptibility as a pickle file.""" 519 assert isinstance(filename, str) 520 if world.rank == 0: 521 with open(filename, 'wb') as fd: 522 pickle.dump((omega_w, G_Gc, chiks_wGG, chi_wGG), fd) 523 524 525def read_component(filename): 526 """Read a stored susceptibility component file""" 527 assert isinstance(filename, str) 528 with open(filename, 'rb') as fd: 529 omega_w, G_Gc, chiks_wGG, chi_wGG = pickle.load(fd) 530 531 return omega_w, G_Gc, chiks_wGG, chi_wGG 532