Source code for mars.tensor.datasource.meshgrid

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from .array import tensor

[docs]def meshgrid(*xi, **kwargs): """ Return coordinate matrices from coordinate vectors. Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector fields over N-D grids, given one-dimensional coordinate tensors x1, x2,..., xn. Parameters ---------- x1, x2,..., xn : array_like 1-D arrays representing the coordinates of a grid. indexing : {'xy', 'ij'}, optional Cartesian ('xy', default) or matrix ('ij') indexing of output. See Notes for more details. sparse : bool, optional If True a sparse grid is returned in order to conserve memory. Default is False. Returns ------- X1, X2,..., XN : Tensor For vectors `x1`, `x2`,..., 'xn' with lengths ``Ni=len(xi)`` , return ``(N1, N2, N3,...Nn)`` shaped tensors if indexing='ij' or ``(N2, N1, N3,...Nn)`` shaped tensors if indexing='xy' with the elements of `xi` repeated to fill the matrix along the first dimension for `x1`, the second for `x2` and so on. Notes ----- This function supports both indexing conventions through the indexing keyword argument. Giving the string 'ij' returns a meshgrid with matrix indexing, while 'xy' returns a meshgrid with Cartesian indexing. In the 2-D case with inputs of length M and N, the outputs are of shape (N, M) for 'xy' indexing and (M, N) for 'ij' indexing. In the 3-D case with inputs of length M, N and P, outputs are of shape (N, M, P) for 'xy' indexing and (M, N, P) for 'ij' indexing. The difference is illustrated by the following code snippet:: xv, yv = mt.meshgrid(x, y, sparse=False, indexing='ij') for i in range(nx): for j in range(ny): # treat xv[i,j], yv[i,j] xv, yv = mt.meshgrid(x, y, sparse=False, indexing='xy') for i in range(nx): for j in range(ny): # treat xv[j,i], yv[j,i] In the 1-D and 0-D case, the indexing and sparse keywords have no effect. Examples -------- >>> import mars.tensor as mt >>> nx, ny = (3, 2) >>> x = mt.linspace(0, 1, nx) >>> y = mt.linspace(0, 1, ny) >>> xv, yv = mt.meshgrid(x, y) >>> xv.execute() array([[ 0. , 0.5, 1. ], [ 0. , 0.5, 1. ]]) >>> yv.execute() array([[ 0., 0., 0.], [ 1., 1., 1.]]) >>> xv, yv = mt.meshgrid(x, y, sparse=True) # make sparse output arrays >>> xv.execute() array([[ 0. , 0.5, 1. ]]) >>> yv.execute() array([[ 0.], [ 1.]]) `meshgrid` is very useful to evaluate functions on a grid. >>> import matplotlib.pyplot as plt >>> x = mt.arange(-5, 5, 0.1) >>> y = mt.arange(-5, 5, 0.1) >>> xx, yy = mt.meshgrid(x, y, sparse=True) >>> z = mt.sin(xx**2 + yy**2) / (xx**2 + yy**2) >>> h = plt.contourf(x,y,z) """ from ..base import broadcast_to indexing = kwargs.pop("indexing", "xy") sparse = kwargs.pop("sparse", False) if kwargs: raise TypeError( f"meshgrid() got an unexpected keyword argument '{list(kwargs)[0]}'" ) if indexing not in ("xy", "ij"): raise ValueError("Valid values for `indexing` are 'xy' and 'ij'.") xi = [tensor(x) for x in xi] xi = [a.ravel() for a in xi] shape = [x.size for x in xi] if indexing == "xy" and len(xi) > 1: xi[0], xi[1] = xi[1], xi[0] shape[0], shape[1] = shape[1], shape[0] grid = [] for i, x in enumerate(xi): slc = [None] * len(shape) slc[i] = slice(None) r = x[tuple(slc)] if not sparse: r = broadcast_to(r, shape) grid.append(r) if indexing == "xy" and len(xi) > 1: grid[0], grid[1] = grid[1], grid[0] return grid