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)