Функция meshgrid

Функция meshgrid создает координатные сетки из одномерных тензоров. Она полезна для генерации всех комбинаций координат по заданным осям. Первый аргумент - тензор координат по оси x (или первому измерению), второй - по оси y (или второму измерению). По умолчанию функция возвращает кортеж из двух тензоров: матрицу координат по первой оси и матрицу координат по второй оси. Параметр indexing управляет порядком индексации.

Синтаксис

torch.meshgrid(x, y, [indexing])

Параметр indexing принимает значения:

'ij' - матричная индексация, при которой первый тензор повторяет строки, второй - столбцы (по умолчанию)
'xy' - декартова индексация, при которой первый тензор повторяет столбцы, второй - строки

Пример

Создадим простую координатную сетку для точек на плоскости:

import torch

x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5])
xx, yy = torch.meshgrid(x, y)
print(xx)
print(yy)

Результат выполнения кода:

tensor([
	[1, 1],
	[2, 2],
	[3, 3],
])
tensor([
	[4, 5],
	[4, 5],
	[4, 5],
])

Здесь матрица xx содержит координаты по оси x, а yy - по оси y.

Пример

Теперь используем декартову индексацию с параметром indexing='xy':

import torch

x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5])
xx, yy = torch.meshgrid(x, y, indexing='xy')
print(xx)
print(yy)

Результат выполнения кода:

tensor([
	[1, 2, 3],
	[1, 2, 3],
])
tensor([
	[4, 4, 4],
	[5, 5, 5],
])

В режиме 'xy' форма сетки соответствует стандартной математической нотации, где первое измерение - это строки (ось y), второе - столбцы (ось x).

Пример

Применим meshgrid для вычисления функции на двумерной области:

import torch

x = torch.linspace(-1, 1, 3)
y = torch.linspace(-1, 1, 3)
xx, yy = torch.meshgrid(x, y, indexing='ij')
z = xx**2 + yy**2
print(z)

Результат выполнения кода:

tensor([
	[2.0000, 2.0000, 2.0000],
	[1.0000, 1.0000, 1.0000],
	[0.0000, 0.0000, 0.0000],
])

Мы вычислили значения функции z = x² + y² в каждой точке сетки.

Смотрите также

функцию linspace,
которая создает линейно разделенные значения на интервале
функцию arange,
которая создает последовательность целых чисел
функцию broadcast_tensors,
которая расширяет формы тензоров по правилам широковещания
функцию stack,
которая объединяет тензоры вдоль новой оси