def __init__(self, ksl, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer, atomic_hamiltonian=None): WaveFunctions.__init__(self, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer) self.ksl = ksl self.S_qMM = None self.T_qMM = None self.P_aqMi = None if atomic_hamiltonian is None: if ksl.using_blacs: atomic_hamiltonian = 'distributed' else: atomic_hamiltonian = 'dense' if isinstance(atomic_hamiltonian, str): atomic_hamiltonian = get_atomic_hamiltonian(atomic_hamiltonian) self.atomic_hamiltonian = atomic_hamiltonian self.timer.start('TCI: Evaluate splines') self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma) self.timer.stop('TCI: Evaluate splines') self.basis_functions = BasisFunctions( gd, [setup.phit_j for setup in setups], kd, dtype=dtype, cut=True)
def __init__(self, ksl, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer, atomic_correction=None, collinear=True): WaveFunctions.__init__(self, gd, nvalence, setups, bd, dtype, collinear, world, kd, kptband_comm, timer) self.ksl = ksl self.S_qMM = None self.T_qMM = None self.P_aqMi = None self.debug_tci = False if atomic_correction is None: if ksl.using_blacs: atomic_correction = 'scipy' else: atomic_correction = 'dense' if isinstance(atomic_correction, str): atomic_correction = get_atomic_correction(atomic_correction) self.atomic_correction = atomic_correction #self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma) with self.timer('TCI: Evaluate splines'): self.tciexpansions = TCIExpansions.new_from_setups(setups) self.basis_functions = BasisFunctions(gd, [setup.phit_j for setup in setups], kd, dtype=dtype, cut=True)
def get_bfs(calc): wfs = calc.wfs bfs = BasisFunctions(wfs.gd, [setup.phit_j for setup in wfs.setups], wfs.kd, cut=True) bfs.set_positions(calc.atoms.get_scaled_positions() % 1.) return bfs
def get_bfs(calc): wfs = calc.wfs bfs = BasisFunctions(wfs.gd, [setup.phit_j for setup in wfs.setups], wfs.kd, cut=True) bfs.set_positions(wfs.spos_ac) return bfs
def initialize(self, density, hamiltonian, spos_ac): if self.kpt_u[0].psit_nG is None: basis_functions = BasisFunctions( self.gd, [setup.phit_j for setup in self.setups], cut=True) if not self.gamma: basis_functions.set_k_points(self.kd.ibzk_qc) basis_functions.set_positions(spos_ac) elif isinstance(self.kpt_u[0].psit_nG, TarFileReference): self.initialize_wave_functions_from_restart_file() if self.kpt_u[0].psit_nG is not None: density.initialize_from_wavefunctions(self) elif density.nt_sG is None: density.initialize_from_atomic_densities(basis_functions) # Initialize GLLB-potential from basis function orbitals if hamiltonian.xc.type == 'GLLB': hamiltonian.xc.initialize_from_atomic_orbitals(basis_functions) else: # XXX??? # We didn't even touch density, but some combinations in paw.set() # will make it necessary to do this for some reason. density.calculate_normalized_charges_and_mix() hamiltonian.update(density) if self.kpt_u[0].psit_nG is None: self.initialize_wave_functions_from_basis_functions( basis_functions, density, hamiltonian, spos_ac)
def get_bfs(calc): wfs = calc.wfs bfs = BasisFunctions(wfs.gd, [setup.phit_j for setup in wfs.setups], wfs.kpt_comm, cut=True) if not wfs.gamma: bfs.set_k_points(wfs.ibzk_qc) bfs.set_positions(calc.atoms.get_scaled_positions() % 1.) return bfs
def get_density(self, atom_indices=None, gridrefinement=2): """Get sum of atomic densities from the given atom list. All atoms are taken if the list is not given.""" all_atoms = self.calculator.get_atoms() if atom_indices is None: atom_indices = range(len(all_atoms)) density = self.calculator.density spos_ac = all_atoms.get_scaled_positions() rank_a = self.finegd.get_ranks_from_positions(spos_ac) density.set_positions(all_atoms.get_scaled_positions(), AtomPartition(self.finegd.comm, rank_a)) # select atoms atoms = [] D_asp = {} rank_a = [] all_D_asp = self.calculator.density.D_asp all_rank_a = self.calculator.density.atom_partition.rank_a for a in atom_indices: if a in all_D_asp: D_asp[len(atoms)] = all_D_asp.get(a) atoms.append(all_atoms[a]) rank_a.append(all_rank_a[a]) atoms = Atoms(atoms, cell=all_atoms.get_cell(), pbc=all_atoms.get_pbc()) spos_ac = atoms.get_scaled_positions() Z_a = atoms.get_atomic_numbers() par = self.calculator.parameters setups = Setups(Z_a, par.setups, par.basis, XC(par.xc), self.calculator.wfs.world) # initialize self.initialize(setups, self.calculator.timer, np.zeros(len(atoms)), False) self.set_mixer(None) # FIXME nparray causes partitionong.py test to fail self.set_positions(spos_ac, AtomPartition(self.gd.comm, rank_a)) self.D_asp = D_asp basis_functions = BasisFunctions( self.gd, [setup.phit_j for setup in self.setups], cut=True) basis_functions.set_positions(spos_ac) self.initialize_from_atomic_densities(basis_functions) aed_sg, gd = self.get_all_electron_density(atoms, gridrefinement) return aed_sg.sum(axis=0), gd
def get_orbitals(calc): """Get LCAO orbitals on 3D grid by lcao_to_grid method.""" bfs_a = [setup.phit_j for setup in calc.wfs.setups] from gpaw.lfc import BasisFunctions bfs = BasisFunctions(calc.wfs.gd, bfs_a, calc.wfs.kd.comm, cut=True) bfs.set_positions(calc.spos_ac) nLCAO = calc.get_number_of_bands() orb_MG = calc.wfs.gd.zeros(nLCAO) C_M = np.identity(nLCAO) bfs.lcao_to_grid(C_M, orb_MG, q=-1) return orb_MG
def get_density(self, atom_indicees=None): """Get sum of atomic densities from the given atom list. All atoms are taken if the list is not given.""" all_atoms = self.calculator.get_atoms() if atom_indicees is None: atom_indicees = range(len(all_atoms)) density = self.calculator.density density.set_positions(all_atoms.get_scaled_positions() % 1.0, self.calculator.wfs.rank_a) # select atoms atoms = [] D_asp = {} rank_a = [] all_D_asp = self.calculator.density.D_asp all_rank_a = self.calculator.density.rank_a for a in atom_indicees: if a in all_D_asp: D_asp[len(atoms)] = all_D_asp.get(a) atoms.append(all_atoms[a]) rank_a.append(all_rank_a[a]) atoms = Atoms(atoms, cell=all_atoms.get_cell()) spos_ac = atoms.get_scaled_positions() % 1.0 Z_a = atoms.get_atomic_numbers() par = self.calculator.input_parameters setups = Setups(Z_a, par.setups, par.basis, par.lmax, XC(par.xc), self.calculator.wfs.world) self.D_asp = D_asp # initialize self.initialize(setups, par.stencils[1], self.calculator.timer, [0] * len(atoms), False) self.set_mixer(None) self.set_positions(spos_ac, rank_a) basis_functions = BasisFunctions( self.gd, [setup.phit_j for setup in self.setups], cut=True) basis_functions.set_positions(spos_ac) self.initialize_from_atomic_densities(basis_functions) aed_sg, gd = Density.get_all_electron_density(self, atoms, gridrefinement=2) return aed_sg[0], gd
def initialize(self, density, hamiltonian, spos_ac): """Initialize wave-functions, density and hamiltonian. Return (nlcao, nrand) tuple with number of bands intialized from LCAO and random numbers, respectively.""" if self.mykpts[0].psit is None: basis_functions = BasisFunctions( self.gd, [setup.phit_j for setup in self.setups], self.kd, dtype=self.dtype, cut=True) basis_functions.set_positions(spos_ac) else: self.initialize_wave_functions_from_restart_file() if self.mykpts[0].psit is not None: density.initialize_from_wavefunctions(self) elif density.nt_sG is None: density.initialize_from_atomic_densities(basis_functions) # Initialize GLLB-potential from basis function orbitals if hamiltonian.xc.type == 'GLLB': hamiltonian.xc.initialize_from_atomic_orbitals(basis_functions) else: # XXX??? # We didn't even touch density, but some combinations in paw.set() # will make it necessary to do this for some reason. density.calculate_normalized_charges_and_mix() hamiltonian.update(density) if self.mykpts[0].psit is None: if 1: # self.collinear: nlcao = self.initialize_wave_functions_from_basis_functions( basis_functions, density, hamiltonian, spos_ac) else: self.random_wave_functions(0) nlcao = 0 nrand = self.bd.nbands - nlcao else: # We got everything from file: nlcao = 0 nrand = 0 return nlcao, nrand
def get_density(self, atom_indices=None, gridrefinement=2): """Get sum of atomic densities from the given atom list. Parameters ---------- atom_indices : list_like All atoms are taken if the list is not given. gridrefinement : 1, 2, 4 Gridrefinement given to get_all_electron_density Returns ------- type spin summed density, grid_descriptor """ all_atoms = self.calculator.get_atoms() if atom_indices is None: atom_indices = range(len(all_atoms)) # select atoms atoms = self.calculator.get_atoms()[atom_indices] spos_ac = atoms.get_scaled_positions() Z_a = atoms.get_atomic_numbers() par = self.calculator.parameters setups = Setups(Z_a, par.setups, par.basis, XC(par.xc), self.calculator.wfs.world) # initialize self.initialize(setups, self.calculator.timer, np.zeros( (len(atoms), 3)), False) self.set_mixer(None) rank_a = self.gd.get_ranks_from_positions(spos_ac) self.set_positions(spos_ac, AtomPartition(self.gd.comm, rank_a)) basis_functions = BasisFunctions( self.gd, [setup.phit_j for setup in self.setups], cut=True) basis_functions.set_positions(spos_ac) self.initialize_from_atomic_densities(basis_functions) aed_sg, gd = self.get_all_electron_density(atoms, gridrefinement) return aed_sg.sum(axis=0), gd
def __init__(self, ksl, gd, nvalence, setups, bd, dtype, world, kd, timer=None): WaveFunctions.__init__(self, gd, nvalence, setups, bd, dtype, world, kd, timer) self.ksl = ksl self.S_qMM = None self.T_qMM = None self.P_aqMi = None self.timer.start('TCI: Evaluate splines') self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma) self.timer.stop('TCI: Evaluate splines') self.basis_functions = BasisFunctions(gd, [setup.phit_j for setup in setups], kd.comm, cut=True) if not kd.gamma: self.basis_functions.set_k_points(kd.ibzk_qc)
def get_all_electron_density(self, atoms=None, gridrefinement=2, spos_ac=None, skip_core=False): """Return real all-electron density array. Usage: Either get_all_electron_density(atoms) or get_all_electron_density(spos_ac=spos_ac) skip_core=True theoretically returns the all-electron valence density (use with care; will not in general integrate to valence) """ if spos_ac is None: spos_ac = atoms.get_scaled_positions() % 1.0 # Refinement of coarse grid, for representation of the AE-density # XXXXXXXXXXXX think about distribution depending on gridrefinement! if gridrefinement == 1: gd = self.redistributor.aux_gd n_sg = self.nt_sG.copy() # This will get the density with the same distribution # as finegd: n_sg = self.redistributor.distribute(n_sg) elif gridrefinement == 2: gd = self.finegd if self.nt_sg is None: self.interpolate_pseudo_density() n_sg = self.nt_sg.copy() elif gridrefinement == 4: # Extra fine grid gd = self.finegd.refine() # Interpolation function for the density: interpolator = Transformer(self.finegd, gd, 3) # XXX grids! # Transfer the pseudo-density to the fine grid: n_sg = gd.empty(self.nspins) if self.nt_sg is None: self.interpolate_pseudo_density() for s in range(self.nspins): interpolator.apply(self.nt_sg[s], n_sg[s]) else: raise NotImplementedError # Add corrections to pseudo-density to get the AE-density splines = {} phi_aj = [] phit_aj = [] nc_a = [] nct_a = [] for a, id in enumerate(self.setups.id_a): if id in splines: phi_j, phit_j, nc, nct = splines[id] else: # Load splines: phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4] splines[id] = (phi_j, phit_j, nc, nct) phi_aj.append(phi_j) phit_aj.append(phit_j) nc_a.append([nc]) nct_a.append([nct]) # Create localized functions from splines phi = BasisFunctions(gd, phi_aj) phit = BasisFunctions(gd, phit_aj) nc = LFC(gd, nc_a) nct = LFC(gd, nct_a) phi.set_positions(spos_ac) phit.set_positions(spos_ac) nc.set_positions(spos_ac) nct.set_positions(spos_ac) I_sa = np.zeros((self.nspins, len(spos_ac))) a_W = np.empty(len(phi.M_W), np.intc) W = 0 for a in phi.atom_indices: nw = len(phi.sphere_a[a].M_w) a_W[W:W + nw] = a W += nw x_W = phi.create_displacement_arrays()[0] D_asp = self.D_asp # XXX really? rho_MM = np.zeros((phi.Mmax, phi.Mmax)) for s, I_a in enumerate(I_sa): M1 = 0 for a, setup in enumerate(self.setups): ni = setup.ni D_sp = D_asp.get(a) if D_sp is None: D_sp = np.empty((self.nspins, ni * (ni + 1) // 2)) else: I_a[a] = ( (setup.Nct) / self.nspins - sqrt(4 * pi) * np.dot(D_sp[s], setup.Delta_pL[:, 0])) if not skip_core: I_a[a] -= setup.Nc / self.nspins if gd.comm.size > 1: gd.comm.broadcast(D_sp, D_asp.partition.rank_a[a]) M2 = M1 + ni rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s]) M1 = M2 assert np.all(n_sg[s].shape == phi.gd.n_c) phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a, x_W) phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a, x_W) a_W = np.empty(len(nc.M_W), np.intc) W = 0 for a in nc.atom_indices: nw = len(nc.sphere_a[a].M_w) a_W[W:W + nw] = a W += nw scale = 1.0 / self.nspins for s, I_a in enumerate(I_sa): if not skip_core: nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a) nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a) gd.comm.sum(I_a) N_c = gd.N_c g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c if not skip_core: for I, g_c in zip(I_a, g_ac): if (g_c >= 0).all() and (g_c < gd.n_c).all(): n_sg[s][tuple(g_c)] -= I / gd.dv return n_sg, gd
def new_get_all_electron_density(self, atoms, gridrefinement=2): """Return real all-electron density array.""" # Refinement of coarse grid, for representation of the AE-density if gridrefinement == 1: gd = self.gd n_sg = self.nt_sG.copy() elif gridrefinement == 2: gd = self.finegd if self.nt_sg is None: self.interpolate() n_sg = self.nt_sg.copy() elif gridrefinement == 4: # Extra fine grid gd = self.finegd.refine() # Interpolation function for the density: interpolator = Transformer(self.finegd, gd, 3) # Transfer the pseudo-density to the fine grid: n_sg = gd.empty(self.nspins) if self.nt_sg is None: self.interpolate() for s in range(self.nspins): interpolator.apply(self.nt_sg[s], n_sg[s]) else: raise NotImplementedError # Add corrections to pseudo-density to get the AE-density splines = {} phi_aj = [] phit_aj = [] nc_a = [] nct_a = [] for a, id in enumerate(self.setups.id_a): if id in splines: phi_j, phit_j, nc, nct = splines[id] else: # Load splines: phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4] splines[id] = (phi_j, phit_j, nc, nct) phi_aj.append(phi_j) phit_aj.append(phit_j) nc_a.append([nc]) nct_a.append([nct]) # Create localized functions from splines phi = BasisFunctions(gd, phi_aj) phit = BasisFunctions(gd, phit_aj) nc = LFC(gd, nc_a) nct = LFC(gd, nct_a) spos_ac = atoms.get_scaled_positions() % 1.0 phi.set_positions(spos_ac) phit.set_positions(spos_ac) nc.set_positions(spos_ac) nct.set_positions(spos_ac) I_sa = np.zeros((self.nspins, len(atoms))) a_W = np.empty(len(phi.M_W), np.int32) W = 0 for a in phi.atom_indices: nw = len(phi.sphere_a[a].M_w) a_W[W:W + nw] = a W += nw rho_MM = np.zeros((phi.Mmax, phi.Mmax)) for s, I_a in enumerate(I_sa): M1 = 0 for a, setup in enumerate(self.setups): ni = setup.ni D_sp = self.D_asp.get(a) if D_sp is None: D_sp = np.empty((self.nspins, ni * (ni + 1) // 2)) else: I_a[a] = ((setup.Nct - setup.Nc) / self.nspins - sqrt(4 * pi) * np.dot(D_sp[s], setup.Delta_pL[:, 0])) if gd.comm.size > 1: gd.comm.broadcast(D_sp, self.rank_a[a]) M2 = M1 + ni rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s]) M1 = M2 phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a) phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a) a_W = np.empty(len(nc.M_W), np.int32) W = 0 for a in nc.atom_indices: nw = len(nc.sphere_a[a].M_w) a_W[W:W + nw] = a W += nw scale = 1.0 / self.nspins for s, I_a in enumerate(I_sa): nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a) nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a) gd.comm.sum(I_a) N_c = gd.N_c g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c for I, g_c in zip(I_a, g_ac): if (g_c >= 0).all() and (g_c < gd.n_c).all(): n_sg[s][tuple(g_c)] -= I / gd.dv return n_sg, gd
# / X : \ # +--/-----------/-\--:--------\----------+ # | | | | : | | # | | | | : | | # | | x | | : x | | # | | | | : | | # | | | | : | | # +--\-----------\-/--:--------/----------+ # \ X : / # \ / \ : / # --------- --:------ # : # # ':' is the domain wall if split on two cpu's gd = GridDescriptor(N_c=[40, 8, 8], cell_cv=[10., 2., 2.], pbc_c=(0, 1, 1)) pos_ac = np.array([[.25, .5, .5], [.55, .5, .5]]) kpts_kc = None spline = Spline(l=0, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0]), r_g=None, beta=None, points=25) spline_aj = [[spline] for pos_c in pos_ac] bfs = BasisFunctions(gd, spline_aj) if kpts_kc is not None: bfs.set_k_points(kpts_kc) bfs.set_positions(pos_ac)
def paw_corrections(self, gridrefinement=2): Fn_wsg, gd = self.interpolate_pseudo_density(gridrefinement) # Splines splines = {} phi_aj = [] phit_aj = [] for a, id in enumerate(self.setups.id_a): if id in splines: phi_j, phit_j = splines[id] else: # Load splines: phi_j, phit_j = self.setups[a].get_partial_waves()[:2] splines[id] = (phi_j, phit_j) phi_aj.append(phi_j) phit_aj.append(phit_j) # Create localized functions from splines phi = BasisFunctions(gd, phi_aj, dtype=float) phit = BasisFunctions(gd, phit_aj, dtype=float) # phi = BasisFunctions(gd, phi_aj, dtype=complex) # phit = BasisFunctions(gd, phit_aj, dtype=complex) spos_ac = self.atoms.get_scaled_positions() phi.set_positions(spos_ac) phit.set_positions(spos_ac) tmp_g = gd.empty(dtype=float) rho_MM = np.zeros((phi.Mmax, phi.Mmax), dtype=self.dtype) rho2_MM = np.zeros_like(rho_MM) for w in range(self.nw): for s in range(self.nspins): rho_MM[:] = 0 M1 = 0 for a, setup in enumerate(self.setups): ni = setup.ni FD_wsp = self.FD_awsp.get(a) if FD_wsp is None: FD_p = np.empty((ni * (ni + 1) // 2), dtype=self.dtype) else: FD_p = FD_wsp[w][s] if gd.comm.size > 1: gd.comm.broadcast(FD_p, self.rank_a[a]) D_ij = unpack2(FD_p) # unpack does complex conjugation that we don't want so # remove conjugation D_ij = np.triu(D_ij, 1) + np.conj(np.tril(D_ij)) # if FD_wsp is None: # FD_wsp = np.empty((self.nw, self.nspins, # ni * (ni + 1) // 2), # dtype=self.dtype) # if gd.comm.size > 1: # gd.comm.broadcast(FD_wsp, self.rank_a[a]) # D_ij = unpack2(FD_wsp[w][s]) # D_ij = np.triu(D_ij, 1) + np.conj(np.tril(D_ij)) M2 = M1 + ni rho_MM[M1:M2, M1:M2] = D_ij M1 = M2 # Add real part of AE corrections tmp_g[:] = 0 rho2_MM[:] = rho_MM.real # TODO: use ae_valence_density_correction phi.construct_density(rho2_MM, tmp_g, q=-1) phit.construct_density(-rho2_MM, tmp_g, q=-1) # phi.lfc.ae_valence_density_correction(rho2_MM, tmp_g, # np.zeros(len(phi.M_W), # np.intc), # np.zeros(self.na)) # phit.lfc.ae_valence_density_correction(-rho2_MM, tmp_g, # np.zeros(len(phi.M_W), # np.intc), # np.zeros(self.na)) Fn_wsg[w][s] += tmp_g # Add imag part of AE corrections tmp_g[:] = 0 rho2_MM[:] = rho_MM.imag # TODO: use ae_valence_density_correction phi.construct_density(rho2_MM, tmp_g, q=-1) phit.construct_density(-rho2_MM, tmp_g, q=-1) # phi.lfc.ae_valence_density_correction(rho2_MM, tmp_g, # np.zeros(len(phi.M_W), # np.intc), # np.zeros(self.na)) # phit.lfc.ae_valence_density_correction(-rho2_MM, tmp_g, # np.zeros(len(phi.M_W), # np.intc), # np.zeros(self.na)) Fn_wsg[w][s] += 1.0j * tmp_g return Fn_wsg, gd
def with_ae_corrections(self, finegrid=False): """Get pair density including the AE corrections""" nij_g = self.get(finegrid) # Generate the density matrix D_ap = {} # D_aii = {} for a, P_ni in self.P_ani.items(): Pi_i = P_ni[self.i] Pj_i = P_ni[self.j] D_ii = np.outer(Pi_i.conj(), Pj_i) # Note: D_ii is not symmetric but the products of partial waves are # so that we can pack D_ap[a] = pack(D_ii) # D_aii[a] = D_ii # Load partial waves if needed if ((finegrid and (not hasattr(self, 'phi'))) or ((not finegrid) and (not hasattr(self, 'Phi')))): # Splines splines = {} phi_aj = [] phit_aj = [] for a, id in enumerate(self.setups.id_a): if id in splines: phi_j, phit_j = splines[id] else: # Load splines: phi_j, phit_j = self.setups[a].get_partial_waves()[:2] splines[id] = (phi_j, phit_j) phi_aj.append(phi_j) phit_aj.append(phit_j) # Store partial waves as class variables if finegrid: gd = self.density.finegd self.__class__.phi = BasisFunctions(gd, phi_aj) self.__class__.phit = BasisFunctions(gd, phit_aj) self.__class__.phi.set_positions(self.spos_ac) self.__class__.phit.set_positions(self.spos_ac) else: gd = self.density.gd self.__class__.Phi = BasisFunctions(gd, phi_aj) self.__class__.Phit = BasisFunctions(gd, phit_aj) self.__class__.Phi.set_positions(self.spos_ac) self.__class__.Phit.set_positions(self.spos_ac) # Add AE corrections if finegrid: phi = self.phi phit = self.phit gd = self.density.finegd else: phi = self.Phi phit = self.Phit gd = self.density.gd rho_MM = np.zeros((phi.Mmax, phi.Mmax)) M1 = 0 for a, setup in enumerate(self.setups): ni = setup.ni D_p = D_ap.get(a) if D_p is None: D_p = np.empty((ni * (ni + 1) // 2)) if gd.comm.size > 1: gd.comm.broadcast(D_p, self.wfs.partition.rank_a[a]) D_ii = unpack2(D_p) # D_ii = D_aii.get(a) # if D_ii is None: # D_ii = np.empty((ni, ni)) # if gd.comm.size > 1: # gd.comm.broadcast(D_ii, self.wfs.atom_partition.rank_a[a]) M2 = M1 + ni rho_MM[M1:M2, M1:M2] = D_ii M1 = M2 # construct_density assumes symmetric rho_MM and # takes only the upper half of it phi.construct_density(rho_MM, nij_g, q=-1) phit.construct_density(-rho_MM, nij_g, q=-1) # TODO: use ae_valence_density_correction # phi.lfc.ae_valence_density_correction( # rho_MM, nij_g, np.zeros(len(phi.M_W), np.intc), np.zeros(self.na)) # phit.lfc.ae_valence_density_correction( # -rho_MM, nij_g, np.zeros(len(phit.M_W), np.intc), np.zeros(self.na)) return nij_g