将网格上的人口转换为坐标，反之亦然

Question

对于生态项目，我需要在方形网格世界中的两种人口表示之间来回切换：

表示 1: 简单的网格（二维 Numpy 数组），其中每个单元格中的值对应于该单元格中的个体数量。例如，使用 3x3 网格：

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

表示 2： 一个二维 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

将网格上的人口转换为坐标，反之亦然

Converting a population on a grid to coordinates, and vice versa

python

numpy

vectorization