Обобщите использование numpy.meshgrid

Question

Ниже приведен минимальный пример создания расширенной сетки с использованием функции meshgrid:

from numpy import array, meshgrid

foo = array(([1, 2, 3], [4, 5]))
array(meshgrid(foo[0], foo[1])).T.reshape(-1, 2)

##  array([[1, 4],
##         [1, 5],
##         [2, 4],
##         [2, 5],
##         [3, 4],
##         [3, 5]])

Я хотел бы реализовать то же самое, но я не хочу жестко кодировать форму foo.

Вот второй минимальный пример с другой формой для foo:

from numpy import array, meshgrid

foo = array(([1, 2, 3], [4, 5], [6, 7]))
array(meshgrid(foo[0], foo[1], foo[2])).T.reshape(-1, 3)

##  array([[1, 4, 6],
##         [1, 5, 6],
##         [2, 4, 6],
##         [2, 5, 6],
##         [3, 4, 6],
##         [3, 5, 6],
##         [1, 4, 7],
##         [1, 5, 7],
##         [2, 4, 7],
##         [2, 5, 7],
##         [3, 4, 7],
##         [3, 5, 7]])

Как можно легко видеть, единственные изменения, которые я сделал:

• Добавить foo[2] в качестве аргумента к функции meshgrid;
• Измените аргументы reshape с (-1, 2) на (-1, 3).

Может ли кто-нибудь помочь мне обобщить моё решение, чтобы я мог использовать одно и то же выражение в обоих случаях вместо внесения описанных изменений?

Answer 1

Вот как вы можете преобразовать свой вызов, чтобы он работал на любом входе 2D:

from numpy import array, meshgrid

def mesh(foo):
  return array(meshgrid(*foo)).T.reshape(-1, foo.shape[0])

print(mesh(array(([1, 2, 3], [4, 5]))))
print(mesh(array(([1, 2, 3], [4, 5], [6, 7]))))