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