Преобразование населения на сетке в координаты и наоборот

Question

Для экологического проекта мне нужно переключаться между двумя представлениями населения в мире с квадратной сеткой:

Представление 1: Просто сетка (2d Numpy array), где значение в каждой ячейке соответствует количеству особей в этой ячейке. Например, с сеткой 3x3:

grid = np.array(
    [[0, 1, 0],
     [0, 3, 1],
     [0, 0, 0]]
)

Представление 2: Массив 2d Numpy с координатами x, y каждого человека в сетке:

coords = np.array(
    [[0, 1],
     [1, 1],
     [1, 1],
     [1, 1],
     [1, 2]]
)

Как видите, когда в ячейке находится более одного человека, ее координаты повторяются. Следовательно, coords имеет форму (Population_size, 2).

Текущие реализации для grid_to_coords() и coords_to_grid() включают циклы for, как вы можете видеть ниже, что значительно замедляет выполнение:

def grid_to_coords(grid):
    non_zero_pos = np.nonzero(grid)
    pop_size = grid.sum(keepdims=False)
    coords = np.zeros((int(pop_size), 2))
    offset = 0
    for i in range(len(non_zero_pos[0])):
        n_in_pos = int(grid[non_zero_pos[0][i], non_zero_pos[1][i]])
        for j in range(n_in_pos):
            coords[i + j + offset] = [non_zero_pos[0][i], non_zero_pos[1][i]]
        offset += j
    return pos

def coords_to_grid(coords, grid_dim):
    grid = np.zeros((grid_dim, grid_dim), dtype=np.int32)
    for x, y in coords:
        # Add a particle to the grid, making sure it is actually on the grid!
        x = max(0, min(x, grid_dim - 1))
        y = max(0, min(y, grid_dim - 1))
        grid[x, y] += 1
    return grid

Мне нужен способ векторизации этих двух функций. Не могли бы вы помочь? Большое спасибо.

Answer 1

import numpy as np

grid = np.array(
    [[0, 1, 0],
     [0, 3, 1],
     [0, 0, 0]]
)

coords = np.array(
    [[0, 1],
     [1, 1],
     [1, 1],
     [1, 1],
     [1, 2]]
)


def grid_to_coords(grid):
    """
    >>> grid_to_coords(grid)
    array([[0, 1],
           [1, 1],
           [1, 1],
           [1, 1],
           [1, 2]])
    """
    x, y = np.nonzero(grid) # x = [0 1 1]; y = [1 1 2]
    # np.c_[x, y] = [[0 1]
    #                [1 1]
    #                [1 2]]
    # grid[x, y] = [1 3 1]
    return np.c_[x, y].repeat(grid[x, y], axis=0)


def coords_to_grid(coords, grid_dim):
    """
    >>> coords_to_grid(coords, 3)
    array([[0, 1, 0],
           [0, 3, 1],
           [0, 0, 0]])
    """
    unique, counts = np.unique(coords, axis=0, return_counts=True)
    # unique = [[0 1]
    #           [1 1]
    #           [1 2]]
    # counts = [1 3 1]
    ret = np.zeros((grid_dim, grid_dim), dtype=int)
    ret[unique[:, 0], unique[:, 1]] = counts
    return ret