在 Julia 中切片和广播多维数组：meshgrid 示例

Posted 2023-02-23

【中文标题】在 Julia 中切片和广播多维数组：meshgrid 示例

【英文标题】：Slicing and broadcasting multidimensional arrays in Julia : meshgrid example

【发布时间】：2015-05-17 00:41:39

【问题描述】：

我最近开始通过编写自组织地图的简单实现来学习 Julia。我希望用户指定地图的大小和维度，这意味着我不能真正使用 for 循环来处理地图数组，因为我事先不知道我需要多少层循环。所以我绝对需要在任意维度的数组上工作的广播和切片函数。

现在，我需要构造一个地图索引数组。假设我的地图由大小为mapsize = (5, 10, 15) 的数组定义，我需要构造一个大小为(3, 5, 10, 15) 的数组indices，其中indices[:, a, b, c] 应该返回[a, b, c]。

我来自 Python/NumPy 背景，其中解决方案已经由特定的“函数”mgrid 给出：

indices = numpy.mgrid[:5, :10, :15]
print indices.shape # gives (3, 5, 10, 15)
print indices[:, 1, 2, 3] gives [1, 2, 3]

没想到Julia一开始就有这样的功能，所以我转向广播。在 NumPy 中，广播基于一组我认为非常清晰和合乎逻辑的规则。您可以在同一表达式中使用不同维度的数组，只要每个维度中的大小匹配或其中一个为 1 ：

(5, 10, 15)   broadcasts to  (5, 10, 15) 
   (10,  1)

(5,  1, 15)   also broadcasts to (5, 10, 15) 
(1, 10,  1)

为了帮助解决这个问题，您还可以使用 numpy.newaxis 或 None 轻松地为您的数组添加新维度：

array = numpy.zeros((5, 15))
array[:,None,:] has shape (5, 1, 15)

这有助于轻松广播数组：

A = numpy.arange(5)
B = numpy.arange(10)
C = numpy.arange(15)
bA, bB, bC = numpy.broadcast_arrays(A[:,None,None], B[None,:,None], C[None,None,:])
bA.shape == bB.shape == bC.shape = (5, 10, 15)

使用它，创建indices 数组相当简单：

indices = numpy.array(numpy.broadcast_arrays(A[:,None,None], B[None,:,None], C[None,None,:]))
(indices == numpy.mgrid[:5,:10,:15]).all() returns True

一般情况当然要复杂一些，但可以使用列表推导和切片来解决：

arrays = [ numpy.arange(i)[tuple([None if m!=n else slice(None) for m in range(len(mapsize))])] for n, i in enumerate(mapsize) ]
indices = numpy.array(numpy.broadcast_arrays(*arrays))

回到朱莉娅。我尝试应用相同的基本原理并最终实现与上面代码的arrays 列表等效。由于复合表达式语法，这最终比 NumPy 对应物更简单：

arrays = [ (idx = ones(Int, length(mapsize)); idx[n] = i;reshape([1:i], tuple(idx...))) for (n,i)=enumerate(mapsize) ]

现在我被困在这里，因为我真的不知道如何将广播应用到我的生成数组列表中...broadcast[!] 函数要求应用函数 f，而我没有有什么。我尝试使用 for 循环尝试强制广播：

indices = Array(Int, tuple(unshift!([i for i=mapsize], length(mapsize))...))
for i=1:length(mapsize)
    A[i] = arrays[i]
end

但这给了我一个错误：ERROR: convert has no method matching convert(::TypeInt64, ::ArrayInt64,3)

我这样做是否正确？我是否忽略了一些重要的事情？任何帮助表示赞赏。

【参考方案1】：

如果你运行的是 julia 0.4，你可以这样做：

julia> function mgrid(mapsize)
           T = typeof(CartesianIndex(mapsize))
           indices = Array(T, mapsize)
           for I in eachindex(indices)
               indices[I] = I
           end
           indices
       end

如果能说出来就更好了

indices = [I for I in CartesianRange(CartesianIndex(mapsize))]

我会调查的 :-)。

【参考方案2】：

Julia 中的广播在很大程度上模仿了 NumPy 中的广播，因此您应该希望发现它或多或少遵循相同的简单规则（不确定在并非所有输入都具有相同维度数时填充维度的方法是否不过是一样的，因为 Julia 数组是列优先的）。

许多有用的东西，比如newaxis 索引和broadcast_arrays，但是（还）没有实现。 （我希望他们会这样做。）还要注意，与 NumPy 相比，Julia 中的索引工作方式略有不同：当您在 NumPy 中省略尾随维度的索引时，其余索引默认为冒号。在 Julia 中，它们可以说是默认为 one。

我不确定您是否真的需要 meshgrid 函数，您想要使用它的大多数事情都可以通过使用带有广播操作的 arrays 数组的原始条目来完成。 meshgrid 在 matlab 中有用的主要原因是它在广播方面很糟糕。

但是使用broadcast! 函数完成您想要做的事情非常简单：

# assume mapsize is a vector with the desired shape, e.g. mapsize = [2,3,4]

N = length(mapsize)
# Your line to create arrays below, with an extra initial dimension on each array
arrays = [ (idx = ones(Int, N+1); idx[n+1] = i;reshape([1:i], tuple(idx...))) for (n,i) in enumerate(mapsize) ]

# Create indices and fill it one coordinate at a time
indices = zeros(Int, tuple(N, mapsize...))
for (i,arr) in enumerate(arrays)
    dest = sub(indices, i, [Colon() for j=1:N]...)
    broadcast!(identity, dest, arr)
end

我必须在arrays 的条目上添加一个初始单例维度，以与indices 的轴对齐（@9​​87654329@ 在这里很有用...）。 然后我遍历每个坐标，在indices 的相关部分创建一个子数组（一个视图），并填充它。 （在 Julia 0.4 中，索引将默认返回子数组，但现在我们必须明确使用 sub。

对broadcast! 的调用仅在输入arr=arrays[i] 上评估标识函数identity(x)=x，广播到输出的形状。为此使用identity 函数不会降低效率； broadcast! 根据给定的函数、参数的数量和结果的维数生成一个专门的函数。

【讨论】：

Julia 有 [CartesianIndex()] 来创建一个新轴，例如来自 NumPy 的 newaxis。

【参考方案3】：

我猜这与 MATLAB meshgrid 功能相同。我从来没有真正考虑过推广到两个以上的维度，所以有点难以理解。

首先，这是我完全通用的版本，这有点疯狂，但我想不出更好的方法来做到这一点，而不为常见维度（例如 2、3）生成代码

function numpy_mgridN(dims...)
    X = Any[zeros(Int,dims...) for d in 1:length(dims)]
    for d in 1:length(dims)
        base_idx = Any[1:nd for nd in dims]
        for i in 1:dims[d]
            cur_idx = copy(base_idx)
            cur_idx[d] = i
            X[d][cur_idx...] = i
        end
    end
    @show X
end
X = numpy_mgridN(3,4,5)
@show X[1][1,2,3]  # 1
@show X[2][1,2,3]  # 2
@show X[3][1,2,3]  # 3

现在，我所说的代码生成是指，对于 2D 情况，您可以简单地做

function numpy_mgrid(dim1,dim2)
    X = [i for i in 1:dim1, j in 1:dim2]
    Y = [j for i in 1:dim1, j in 1:dim2]
    return X,Y
end

对于 3D 情况：

function numpy_mgrid(dim1,dim2,dim3)
    X = [i for i in 1:dim1, j in 1:dim2, k in 1:dim3]
    Y = [j for i in 1:dim1, j in 1:dim2, k in 1:dim3]
    Z = [k for i in 1:dim1, j in 1:dim2, k in 1:dim3]
    return X,Y,Z
end

测试，例如

X,Y,Z=numpy_mgrid(3,4,5)
@show X
@show Y
@show Z

我猜mgrid 将它们全部推到一个张量中，所以你可以这样做

all = cat(4,X,Y,Z)

还是有些不同的：

julia> all[1,2,3,:]
1x1x1x3 ArrayInt64,4:
[:, :, 1, 1] =
 1

[:, :, 1, 2] =
 2

[:, :, 1, 3] =
 3

julia> vec(all[1,2,3,:])
3-element ArrayInt64,1:
 1
 2
 3

【讨论】：

太棒了！我稍微修改了代码以获得单张量结果：第二行：X = Array(Int, tuple(unshift!([i for i=dims], length(dims))...))；第八行：X[d, cur_idx...] = i