將網格上的人口轉換為坐標，反之亦然

Question

對於一個生態項目，我需要在方格世界上的兩種人口表示之間來回切換：

表示 1：簡單的網格（一個二維 Numpy 數組），其中每個單元格中的值對應於該單元格中的個體數量。 例如，使用 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]]
)

如您所見，當一個單元格上有超過 1 個個體時，其坐標會重復。 因此， 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