Numpy对2D数组进行操作以生成3D数组

Question

我已经生成了一个(x, y)值的numpy数组作为N x N网格。

grid = np.meshgrid(np.linspace(0, 1, 50), np.linspace(0, 1, 50))[0]
grid.shape  // (50, 50, 1)

我有一个带有两个参数并返回3个值的函数。 即(x, y) -> (a, b, c)

如何在2d numpy数组上应用函数以获得3d numpy数组？

Answer 1

如果您的函数确实需要两个参数，则可能不希望将2d映射到3d，而是将2xMxN映射到3xMxN。 为此，将第一行更改为

gridx, gridy = np.meshgrid(np.linspace(0, 1, 50), np.linspace(0, 1, 50))

甚至使用更经济的ix_ ，它的优点是不交换轴

gridy, gridx = np.ix_(np.linspace(0, 1, 50), np.linspace(0, 1, 50))

如果函数f不处理数组参数，则正如@Jacques Gaudin指出的那样， np.vectorize可能就是您想要的。 请注意， vectorize主要是一种便利功能，它不会使事情变得更快。 它会做一些有用的事情，例如广播，这就是为什么使用ix_实际ix_原因

f_wrapped = np.vectorize(f)
result = f_wrapped(gridy, gridx)

请注意，根据您的情况， result是3个元组，其中包含50 x 50的数组，即按输出分组。 如果要链接矢量化函数，这很方便。 如果要全部集中在一个大数组中，只需将result转换为array并可以选择使用transpose来重新排列轴，例如

 triplets_last = np.array(result).transpose((1, 2, 0))

Answer 2

如果我理解正确，那么您将使用np.vectorize装饰器。 通过使用它，您可以将函数应用于meshgrid 。 您的函数应该只采用一个参数，因为您不传递坐标，而是传递坐标处的值（除非值是带有两个元素的tulpes）。

import numpy as np

grid = np.meshgrid(np.linspace(0, 1, 5), np.linspace(0, 1, 5))[0]

@np.vectorize
def func(a):
    return (a, a**.5, a**2)

res = np.array(list(func(grid)))
print(res.shape)
print(res)