original_docstring: 
          Gathers all SNAP descriptors on rank 0 and sorts them.
  
          This is useful for e.g. parallel preprocessing.
  
          Parameters
          ----------
          snap_descriptors_np : numpy.array
              Numpy array with the SNAP descriptors of this ranks local grid.
  
          use_pickled_comm : bool
              If True, the pickled communication route from mpi4py is used.
              If False, a Recv/Sendv combination is used. I am not entirely
              sure what is faster. Technically Recv/Sendv should be faster,
              but I doubt my implementation is all that optimal. For the pickled
              route we can use gather(), which should be fairly quick.
              However, for large grids, one CANNOT use the pickled route;
              too large python objects will break it. Therefore, I am setting
              the Recv/Sendv route as default.
          
docstring: Gathers all SNAP descriptors on rank 0 and sorts them.
  This is useful for e.g. parallel preprocessing.
  
  Parameters
  
  snap_descriptors_np : numpy.array
  Numpy array with the SNAP descriptors of this ranks local grid.
  
  use_pickled_comm : bool
  If True, the pickled communication route from mpi4py is used.
  If False, a Recv/Sendv combination is used. I am not entirely
  sure what is faster. Technically Recv/Sendv should be faster,
  but I doubt my implementation is all that optimal. For the pickled
  route we can use gather(), which should be fairly quick.
  However, for large grids, one CANNOT use the pickled route;
  too large python objects will break it. Therefore, I am setting
  the Recv/Sendv route as default.
docstring_tokens: Gathers all SNAP descriptors on rank 0 and sorts them . This is useful for e . g . parallel preprocessing . Parameters snap_descriptors_np : numpy . array Numpy array with the SNAP descriptors of this ranks local grid . use_pickled_comm : bool If True the pickled communication route from mpi4py is used . If False a Recv / Sendv combination is used . I am not entirely sure what is faster . Technically Recv / Sendv should be faster but I doubt my implementation is all that optimal . For the pickled route we can use gather () which should be fairly quick . However for large grids one CANNOT use the pickled route ; too large python objects will break it . Therefore I am setting the Recv / Sendv route as default .
code: def gather_descriptors(self, snap_descriptors_np, use_pickled_comm=False):
          """
          Gathers all SNAP descriptors on rank 0 and sorts them.
  
          This is useful for e.g. parallel preprocessing.
  
          Parameters
          ----------
          snap_descriptors_np : numpy.array
              Numpy array with the SNAP descriptors of this ranks local grid.
  
          use_pickled_comm : bool
              If True, the pickled communication route from mpi4py is used.
              If False, a Recv/Sendv combination is used. I am not entirely
              sure what is faster. Technically Recv/Sendv should be faster,
              but I doubt my implementation is all that optimal. For the pickled
              route we can use gather(), which should be fairly quick.
              However, for large grids, one CANNOT use the pickled route;
              too large python objects will break it. Therefore, I am setting
              the Recv/Sendv route as default.
          """
          # Barrier to make sure all ranks have descriptors..
          comm = get_comm()
          comm.Barrier()
  
          # Gather the SNAP descriptors into a list.
          if use_pickled_comm:
              all_snap_descriptors_list = comm.gather(snap_descriptors_np, root=0)
          else:
              sendcounts = np.array(comm.gather(np.shape(snap_descriptors_np)[0],
                                                root=0))
              raw_feature_length = self.fingerprint_length+3
  
              if get_rank() == 0:
                  # print("sendcounts: {}, total: {}".format(sendcounts,
                  #                                          sum(sendcounts)))
  
                  # Preparing the list of buffers.
                  all_snap_descriptors_list = []
                  for i in range(0, get_size()):
                      all_snap_descriptors_list.append(np.empty(sendcounts[i] *
                                                                raw_feature_length,
                                                                dtype=np.float64))
  
                  # No MPI necessary for first rank. For all the others,
                  # collect the buffers.
                  all_snap_descriptors_list[0] = snap_descriptors_np
                  for i in range(1, get_size()):
                      comm.Recv(all_snap_descriptors_list[i], source=i,
                                tag=100+i)
                      all_snap_descriptors_list[i] = \
                          np.reshape(all_snap_descriptors_list[i],
                                     (sendcounts[i], raw_feature_length))
              else:
                  comm.Send(snap_descriptors_np, dest=0, tag=get_rank()+100)
              comm.Barrier()
  
          # if get_rank() == 0:
          #     printout(np.shape(all_snap_descriptors_list[0]))
          #     printout(np.shape(all_snap_descriptors_list[1]))
          #     printout(np.shape(all_snap_descriptors_list[2]))
          #     printout(np.shape(all_snap_descriptors_list[3]))
  
          # Dummy for the other ranks.
          # (For now, might later simply broadcast to other ranks).
          snap_descriptors_full = np.zeros([1,1,1,1])
  
          # Reorder the list.
          if get_rank() == 0:
              # Prepare the SNAP descriptor array.
              nx = self.grid_dimensions[0]
              ny = self.grid_dimensions[1]
              nz = self.grid_dimensions[2]
              snap_descriptors_full = np.zeros(
                  [nx, ny, nz, self.fingerprint_length])
              # Fill the full SNAP descriptors array.
              for idx, local_snap_grid in enumerate(all_snap_descriptors_list):
                  # We glue the individual cells back together, and transpose.
                  first_x = int(local_snap_grid[0][0])
                  first_y = int(local_snap_grid[0][1])
                  first_z = int(local_snap_grid[0][2])
                  last_x = int(local_snap_grid[-1][0])+1
                  last_y = int(local_snap_grid[-1][1])+1
                  last_z = int(local_snap_grid[-1][2])+1
                  snap_descriptors_full[first_x:last_x,
                                        first_y:last_y,
                                        first_z:last_z] = \
                      np.reshape(local_snap_grid[:,3:],[last_z-first_z,
                                                        last_y-first_y,
                                                        last_x-first_x,
                                                        self.fingerprint_length]).transpose([2, 1, 0, 3])
  
                  # Leaving this in here for debugging purposes.
                  # This is the slow way to reshape the descriptors.
                  # for entry in local_snap_grid:
                  #     x = int(entry[0])
                  #     y = int(entry[1])
                  #     z = int(entry[2])
                  #     snap_descriptors_full[x, y, z] = entry[3:]
          return snap_descriptors_full
code_tokens: def gather_descriptors ( self , snap_descriptors_np , use_pickled_comm = False ) : comm = get_comm ( ) comm . Barrier ( ) if use_pickled_comm : all_snap_descriptors_list = comm . gather ( snap_descriptors_np , root = 0 ) else : sendcounts = np . array ( comm . gather ( np . shape ( snap_descriptors_np ) [ 0 ] , root = 0 ) ) raw_feature_length = self . fingerprint_length + 3 if get_rank ( ) == 0 : all_snap_descriptors_list = [ ] for i in range ( 0 , get_size ( ) ) : all_snap_descriptors_list . append ( np . empty ( sendcounts [ i ] * raw_feature_length , dtype = np . float64 ) ) all_snap_descriptors_list [ 0 ] = snap_descriptors_np for i in range ( 1 , get_size ( ) ) : comm . Recv ( all_snap_descriptors_list [ i ] , source = i , tag = 100 + i ) all_snap_descriptors_list [ i ] = np . reshape ( all_snap_descriptors_list [ i ] , ( sendcounts [ i ] , raw_feature_length ) ) else : comm . Send ( snap_descriptors_np , dest = 0 , tag = get_rank ( ) + 100 ) comm . Barrier ( ) snap_descriptors_full = np . zeros ( [ 1 , 1 , 1 , 1 ] ) if get_rank ( ) == 0 : nx = self . grid_dimensions [ 0 ] ny = self . grid_dimensions [ 1 ] nz = self . grid_dimensions [ 2 ] snap_descriptors_full = np . zeros ( [ nx , ny , nz , self . fingerprint_length ] ) for idx , local_snap_grid in enumerate ( all_snap_descriptors_list ) : first_x = int ( local_snap_grid [ 0 ] [ 0 ] ) first_y = int ( local_snap_grid [ 0 ] [ 1 ] ) first_z = int ( local_snap_grid [ 0 ] [ 2 ] ) last_x = int ( local_snap_grid [ - 1 ] [ 0 ] ) + 1 last_y = int ( local_snap_grid [ - 1 ] [ 1 ] ) + 1 last_z = int ( local_snap_grid [ - 1 ] [ 2 ] ) + 1 snap_descriptors_full [ first_x : last_x , first_y : last_y , first_z : last_z ] = np . reshape ( local_snap_grid [ : , 3 : ] , [ last_z - first_z , last_y - first_y , last_x - first_x , self . fingerprint_length ] ) . transpose ( [ 2 , 1 , 0 , 3 ] ) return snap_descriptors_full
short_docstring: Gathers all SNAP descriptors on rank 0 and sorts them.
short_docstring_tokens: Gathers all SNAP descriptors on rank 0 and sorts them .
comment: """
          Gathers all SNAP descriptors on rank 0 and sorts them.
  
          This is useful for e.g. parallel preprocessing.
  
          Parameters
          ----------
          snap_descriptors_np : numpy.array
              Numpy array with the SNAP descriptors of this ranks local grid.
  
          use_pickled_comm : bool
              If True, the pickled communication route from mpi4py is used.
              If False, a Recv/Sendv combination is used. I am not entirely
              sure what is faster. Technically Recv/Sendv should be faster,
              but I doubt my implementation is all that optimal. For the pickled
              route we can use gather(), which should be fairly quick.
              However, for large grids, one CANNOT use the pickled route;
              too large python objects will break it. Therefore, I am setting
              the Recv/Sendv route as default.
          """$SEP$# Barrier to make sure all ranks have descriptors..$SEP$# Gather the SNAP descriptors into a list.$SEP$# print("sendcounts: {}, total: {}".format(sendcounts,$SEP$#                                          sum(sendcounts)))$SEP$# Preparing the list of buffers.$SEP$# No MPI necessary for first rank. For all the others,$SEP$# collect the buffers.$SEP$# if get_rank() == 0:$SEP$#     printout(np.shape(all_snap_descriptors_list[0]))$SEP$#     printout(np.shape(all_snap_descriptors_list[1]))$SEP$#     printout(np.shape(all_snap_descriptors_list[2]))$SEP$#     printout(np.shape(all_snap_descriptors_list[3]))$SEP$# Dummy for the other ranks.$SEP$# (For now, might later simply broadcast to other ranks).$SEP$# Reorder the list.$SEP$# Prepare the SNAP descriptor array.$SEP$# Fill the full SNAP descriptors array.$SEP$# We glue the individual cells back together, and transpose.$SEP$# Leaving this in here for debugging purposes.$SEP$# This is the slow way to reshape the descriptors.$SEP$# for entry in local_snap_grid:$SEP$#     x = int(entry[0])$SEP$#     y = int(entry[1])$SEP$#     z = int(entry[2])$SEP$#     snap_descriptors_full[x, y, z] = entry[3:]