lsdm.py 文件源码

def initialize(self, comm, config, args):

        rank = comm.Get_rank()
        size = comm.Get_size()

        self.config = config
        self.args = args

        filename = args.struct_file[0] 
        self.struct_filename = filename
        self.npoints,self.natoms = coord_reader.get_nframes_natoms(filename)

        if coord_reader.supports_parallel_reading(filename): 
            # read coordinates in parallel
            self.idxs_thread, self.npoints_per_thread, self.offsets_per_thread = p_index.get_idxs_thread(comm, self.npoints)
            coords_thread = coord_reader.get_coordinates(filename, idxs=self.idxs_thread)
            coords_ravel = coords_thread.ravel()
            ravel_lengths, ravel_offsets = p_index.get_ravel_offsets(self.npoints_per_thread,self.natoms)
            coordstemp = np.zeros(self.npoints*3*self.natoms, dtype='float')
            start = MPI.Wtime()
            comm.Allgatherv(coords_ravel, (coordstemp, ravel_lengths, ravel_offsets, MPI.DOUBLE))
            self.coords = coordstemp.reshape((self.npoints,3,self.natoms))
        else: 
            # serial reading
            if rank == 0:
                self.coords = coord_reader.get_coordinates(filename)
            else:
                self.coords = np.zeros((self.npoints,3,self.natoms),dtype=np.double)
            comm.Bcast(self.coords, root=0) 

        logging.info('input coordinates loaded')

        self.initialize_local_scale()
        self.initialize_weights()
        self.initialize_metric()