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

Question

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

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

我来自 Python/NumPy 背景，其中解决方案已由特定 "function"、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(::Type{Int64}, ::Array{Int64,3})

我这样做对吗？我是否忽略了一些重要的事情？感谢您的帮助。

Answer 1

我猜这与 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 Array{Int64,4}:
[:, :, 1, 1] =
 1

[:, :, 1, 2] =
 2

[:, :, 1, 3] =
 3

julia> vec(all[1,2,3,:])
3-element Array{Int64,1}:
 1
 2
 3

Answer 2

如果你是运行 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))]

我会调查一下:-)。

Answer 3

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

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

我不确定您是否真的需要 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 的轴对齐（newaxis 在这里很有用......）。然后我遍历每个坐标，在 indices 的相关部分创建一个子数组（一个视图）并填充它。（索引将默认返回 Julia 0.4 中的子数组，但现在我们必须显式使用 sub）。

对 broadcast! 的调用只是计算输入 arr=arrays[i] 上的身份函数 identity(x)=x，广播到输出的形状。为此使用 identity 函数不会降低效率； broadcast! 根据给定函数、参数数量和结果的维数生成专用函数。

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

Slicing and broadcasting multidimensional arrays in Julia : meshgrid example

slice

multidimensional-array

julia