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
