Как можно векторизовать функцию для двух, взять в качестве аргументов две последовательности разной длины (m и n) и вернуть матрицу (mxn) в Python NumPy?

Question

У меня есть последовательности (например) и функция:

import numpy as np

a = np.arange(10,21)
b = np.arange(10,16)

def f(x,y):
    return x**2+y*3

, и я хочу построить матрицу (10x5), чтобы каждый из ее элементов был результатом функции на эти элементы итеративно. Так, например, элемент строки 2nd и столбца 3rd матрицы будет:

matrix[1,2]
>>> 157

, а именно f (11,12) = 11 ** 2 + 12 * 3 = 121 + 36 = 157.

Я пытался использовать вложенные циклы (два цикла "for"), но это слишком медленно. Есть ли способ для numpy выполнить операции сразу, например, функция np, которая принимает в качестве аргумента a, b и f и возвращает желаемый результат?

Большое спасибо: )

Answer 1

Вы можете превратить a в вектор-столбец и вуаля:

matrix = a[:,np.newaxis] ** 2 + b * 3

Хотя что-то вроде этого чище:

a, b = np.ogrid[10:21,10:16]
matrix = a ** 2 + b * 3

(отредактировано благодаря комментарию hpaulj)