diff --git a/cunumeric/module.py b/cunumeric/module.py index 5981e06d7..c579e8e50 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -789,6 +789,110 @@ def linspace( return y.astype(dtype, copy=False) +# meshgrid is adapted from the numpy implementation. +def meshgrid( + *xii: Any, copy: bool = True, sparse: bool = False, indexing: str = "xy" +) -> list[ndarray]: + """ + 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 arrays 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 the shape of the returned coordinate array for dimension *i* + is reduced from ``(N1, ..., Ni, ... Nn)`` to + ``(1, ..., 1, Ni, 1, ..., 1)``. These sparse coordinate grids are + intended to be use with :ref:`basics.broadcasting`. When all + coordinates are used in an expression, broadcasting still leads to a + fully-dimensonal result array. + + Default is False. + copy : bool, optional + If False, a view into the original arrays are returned in order to + conserve memory. Default is True. Please note that + ``sparse=False, copy=False`` will likely return non-contiguous + arrays. Furthermore, more than one element of a broadcast array + may refer to a single memory location. If you need to write to the + arrays, make copies first. + + Returns + ------- + X1, X2,..., XN : ndarray + For vectors `x1`, `x2`,..., 'xn' with lengths ``Ni=len(xi)`` , + return ``(N1, N2, N3,...Nn)`` shaped arrays if indexing='ij' + or ``(N2, N1, N3,...Nn)`` shaped arrays 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 = np.meshgrid(x, y, indexing='ij') + for i in range(nx): + for j in range(ny): + # treat xv[i,j], yv[i,j] + + xv, yv = np.meshgrid(x, y, 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. + + See Also + -------- + mgrid : Construct a multi-dimensional "meshgrid" using indexing notation. + ogrid : Construct an open multi-dimensional "meshgrid" using indexing + notation. + + Availability + -------- + Multiple GPUs, Multiple CPUs + """ + xi = [convert_to_cunumeric_ndarray(x) for x in xii] + ndim = len(xi) + + if indexing not in ["xy", "ij"]: + raise ValueError("Valid values for `indexing` are 'xy' and 'ij'.") + + s0 = (1,) * ndim + output = [ + np.asanyarray(x).reshape(s0[:i] + (-1,) + s0[i + 1 :]) + for i, x in enumerate(xi) + ] + + if indexing == "xy" and ndim > 1: + # Switch first and second axis. + output[0].shape = (1, -1) + s0[2:] + output[1].shape = (-1, 1) + s0[2:] + + if not sparse: + # Return the full N-D matrix (not only the 1-D vector). + output = np.broadcast_arrays(*output, subok=True) + + if copy: + output = [x.copy() for x in output] + + return output + + # Building matrices diff --git a/tests/integration/test_linspace.py b/tests/integration/test_linspace.py index 170bec59c..107a253fb 100644 --- a/tests/integration/test_linspace.py +++ b/tests/integration/test_linspace.py @@ -57,6 +57,21 @@ def test_axis(): assert np.array_equal(z, w) +@pytest.mark.parametrize("sparse", [True, False]) +@pytest.mark.parametrize("indexing", ["xy", "ij"]) +def test_meshgrid(sparse, indexing): + xnp = np.linspace(0.0, 1.0, 10) + ynp = np.linspace(0.5, 1.5, 10) + Xnp, Ynp = np.meshgrid(xnp, ynp, sparse=sparse, indexing=indexing) + + xnum = num.linspace(0.0, 1.0, 10) + ynum = num.linspace(0.5, 1.5, 10) + Xnum, Ynum = num.meshgrid(xnum, ynum, sparse=sparse, indexing=indexing) + + assert num.array_equal(Xnum, Xnp) + assert num.array_equal(Ynum, Ynp) + + if __name__ == "__main__": import sys