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)
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, 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)
