def windows_iteration(shape, max_bunch, l, center, excluded): """Main iteration over all the grid by number of neighs. Parameters ---------- shape: tuple the size of the grid. max_bunch: int maximum number of elements we want to retrieve their neighbours at the same time. l: int size of the windows to retrieve. center: int the place of the center. excluded: boolean (default = False) if we want to exclude the element `i`. Returns ------- inds: np.ndarray the indices of the elements which we want to retrieve their neighbours. neighs: np.ndarray the neighbours of each inds. rel_pos: np.ndarray the relative positions of each neighbours to its retrieved element. """ ## Splitting by special corner and border points shapes = np.array(list(np.cumprod(shape[1:][::-1])[::-1]) + [1]) coord2ind = create_map2indices(shape) borders, ranges, sizes_nei =\ get_irregular_indices_grid(shape, l, center, excluded) ## Retrieving by sizes of neighbourhood sizes_u = np.unique(sizes_nei) if len(sizes_u) != 1: for s in sizes_u: inds, neighs = [], [] size_idx = np.where(s == sizes_nei)[0] for idx_s in size_idx: # Get indices of the special points inds_s = get_indices_from_borders(borders[idx_s], coord2ind) # Get relative position rel_pos = get_relative_neighs_comb(ranges[idx_s]) # Get neigbours neighs_s = get_irregular_neighsmatrix(inds_s, rel_pos, shapes) inds.append(inds_s) neighs.append(neighs_s) ## Formatting properly inds, neighs = np.concatenate(inds), np.concatenate(neighs, axis=1) inds, neighs = inds.astype(int), neighs.astype(int) yield inds, neighs, rel_pos indices, relative_neighs, rel_pos =\ get_indices_constant_regular(shape, coord2ind, l, center, excluded) indices_split = split_parallel(indices, max_bunch) # indices_split = [indices] for inds in indices_split: neighs = get_regular_neighsmatrix(inds, relative_neighs) yield inds, neighs, rel_pos