class UTDensityFunctionSetup(UTLocalizedFunctionSetup): __doc__ = UTLocalizedFunctionSetup.__doc__ + """ Atomic density matrices are distributed over domains.""" def setUp(self): UTLocalizedFunctionSetup.setUp(self) self.finegd = self.gd.refine() self.density = Density(self.gd, self.finegd, self.nspins, p.charge) self.density.initialize(self.setups, p.stencils[1], self.timer, \ self.atoms.get_initial_magnetic_moments(), p.hund) self.density.D_asp = {} self.density.rank_a = self.rank0_a self.allocate(self.density.D_asp, self.density.rank_a) assert self.allocated for a, D_sp in self.density.D_asp.items(): ni = self.setups[a].ni for s, D_p in enumerate(D_sp): D_p[:] = 1e9 * self.kd_old.comm.rank + 1e6 * self.bd.comm.rank \ + 1e3 * s + np.arange(ni * (ni + 1) // 2, dtype=float) def tearDown(self): del self.density UTLocalizedFunctionSetup.tearDown(self) # ================================= def test_initial_consistency(self): self.update_references(self.density.D_asp, self.density.rank_a) self.check_values(self.density.D_asp, self.density.rank_a) def test_redistribution_to_domains(self): self.update_references(self.density.D_asp, self.density.rank_a) spos_ac = self.atoms.get_scaled_positions() % 1.0 rank_a = self.gd.get_ranks_from_positions(spos_ac) self.density.set_positions(spos_ac, rank_a) self.check_values(self.density.D_asp, self.density.rank_a) def test_redistribution_to_same(self): self.update_references(self.density.D_asp, self.density.rank_a) spos_ac = self.atoms.get_scaled_positions() % 1.0 rank_a = self.gd.get_ranks_from_positions(spos_ac) self.density.set_positions(spos_ac, rank_a) self.density.set_positions(spos_ac, rank_a) self.check_values(self.density.D_asp, self.density.rank_a)
def set_positions(self, spos_ac, rank_a=None): Density.set_positions(self, spos_ac, rank_a) self.nct_q = self.pd2.zeros() self.nct.add(self.nct_q, 1.0 / self.nspins) self.nct_G = self.pd2.ifft(self.nct_q)
def set_positions(self, spos_ac, atom_partition): Density.set_positions(self, spos_ac, atom_partition) self.nct_q = self.pd2.zeros() self.nct.add(self.nct_q, 1.0 / self.nspins) self.nct_G = self.pd2.ifft(self.nct_q)