cunumeric/module: implement numpy.meshgrid

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
 
 

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