def write(self, writer, write_wave_functions=False):
FDPWWaveFunctions.write(self, writer)
if not write_wave_functions:
return
writer.add_array('values',
(self.nspins, self.kd.nibzkpts, self.bd.nbands) +
tuple(self.gd.get_size_of_global_array()), self.dtype)
for s in range(self.nspins):
for k in range(self.kd.nibzkpts):
for n in range(self.bd.nbands):
psit_G = self.get_wave_function_array(n, k, s)
writer.fill(psit_G * Bohr**-1.5)
def write(self, writer, write_wave_functions=False):
FDPWWaveFunctions.write(self, writer)
if not write_wave_functions:
return
writer.add_array(
'coefficients',
(self.nspins, self.kd.nibzkpts, self.bd.nbands, self.pd.ngmax),
complex)
c = units.Bohr**-1.5
for s in range(self.nspins):
for k in range(self.kd.nibzkpts):
for n in range(self.bd.nbands):
psit_G = self.get_wave_function_array(n, k, s,
realspace=False,
cut=False)
writer.fill(psit_G * c)
writer.add_array('indices', (self.kd.nibzkpts, self.pd.ngmax),
np.int32)
if self.bd.comm.rank > 0:
return
Q_G = np.empty(self.pd.ngmax, np.int32)
kk = 0
for r in range(self.kd.comm.size):
for q, ks in enumerate(self.kd.get_indices(r)):
s, k = divmod(ks, self.kd.nibzkpts)
ng = self.ng_k[k]
if s == 1:
return
if r == self.kd.comm.rank:
Q_G[:ng] = self.pd.Q_qG[q]
if r > 0:
self.kd.comm.send(Q_G, 0)
if self.kd.comm.rank == 0:
if r > 0:
self.kd.comm.receive(Q_G, r)
Q_G[ng:] = -1
writer.fill(Q_G)
assert k == kk
kk += 1