Numpy работает над 2D-массивом для создания 3D-массива

Question

Я создал массив numpy значений (x, y) как сетку N x N.

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

У меня есть функция, которая принимает два параметра и возвращает 3 значения. i.e. (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 не обрабатывает аргументы массива затем @ Жак Годин указывает np.vectorize, вероятно, то, что вы хотите. Будьте предупреждены, что vectorize - это прежде всего функция удобства, которая не ускоряет работу. Это действительно полезные вещи, как вещание из-за чего с помощью ix_ на самом деле работает

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

Обратите внимание, что result в вашем случае является 3-кортежем из 50 х 50 массивов, т.е. сгруппирован по выходу. Это удобно, если вы хотите связать векторизованные функции. Если вы хотите, чтобы все в одном большом массиве только конвертировали result в array и, необязательно, использовали transpose для изменения осей, например.

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

Answer 2

Если я правильно понял, вы после декоратора np.vectorize. Используя его, вы можете применить функцию через meshgrid. Ваша функция должна принимать только один параметр, хотя, поскольку вы не передаете координаты, а значение в координатах (если только значения не являются тюльпанами с двумя элементами).

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)