def setUp(self):
UTLocalizedFunctionSetup.setUp(self)
fdksl = get_KohnSham_layouts(None, 'fd', self.gd, self.bd, self.dtype)
lcaoksl = get_KohnSham_layouts(None,
'lcao',
self.gd,
self.bd,
self.dtype,
nao=self.setups.nao)
args = (self.gd, self.setups.nvalence, self.setups, self.bd,
self.dtype, world, self.kd)
self.wfs = FDWaveFunctions(p.stencils[0], fdksl, fdksl, lcaoksl, *args)
self.wfs.rank_a = self.rank0_a
self.allocate(self.wfs.kpt_u, self.wfs.rank_a)
assert self.allocated
for kpt in self.wfs.kpt_u:
for a, P_ni in kpt.P_ani.items():
for myn, P_i in enumerate(P_ni):
n = self.bd.global_index(myn)
P_i[:] = 1e12 * kpt.s + 1e9 * kpt.k + 1e6 * a + 1e3 * n \
+ np.arange(self.setups[a].ni, dtype=self.dtype)
