Adding meshgrid

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -898,12 +898,16 @@ def validate_inputs(inputs: dict, out=None, reduce=False) -> tuple:  # noqa: C90
             fast_path = False
         if first_input.blocks != out.blocks:
             fast_path = False
+        if 0 in out.chunks:  # fast_eval has zero division error for 0 shapes
+            fast_path = False
     # Then, the rest of the operands
     for input_ in NDinputs:
         if first_input.chunks != input_.chunks:
             fast_path = False
         if first_input.blocks != input_.blocks:
             fast_path = False
+        if 0 in input_.chunks:  # fast_eval has zero division error for 0 shapes
+            fast_path = False
 
     return first_input.shape, first_input.chunks, first_input.blocks, fast_path
 
@@ -2269,7 +2273,7 @@ def get_chunk(self, nchunk):
         return out.schunk.get_chunk(nchunk)
 
     def update_expr(self, new_op):  # noqa: C901
-        prev_flag = getattr(blosc2, "_disable_overloaded_equal", False)
+        prev_flag = blosc2._disable_overloaded_equal
         # We use a lot of the original NDArray.__eq__ as 'is', so deactivate the overloaded one
         blosc2._disable_overloaded_equal = True
         # One of the two operands are LazyExpr instances

src/blosc2/ndarray.py

@@ -385,7 +385,7 @@ def _check_allowed_dtypes(
     ):
         raise RuntimeError(
             "Expected LazyExpr, NDArray, NDField, C2Array, Proxy, np.ndarray or scalar instances"
-            f" and you provided a '{type(value)}' instance"
+            + f" and you provided a '{type(value)}' instance"
         )
 
 
@@ -5195,7 +5195,7 @@ def broadcast_to(arr, shape):
     return (arr + blosc2.zeros(shape, dtype=arr.dtype)).compute()  # return lazyexpr quickly
 
 
-def meshgrid(arrays: NDArray, indexing: str = "xy") -> Sequence[NDArray]:
+def meshgrid(*arrays: NDArray, indexing: str = "xy") -> Sequence[NDArray]:
     """
     Returns coordinate matrices from coordinate vectors.
 
@@ -5217,4 +5217,32 @@ def meshgrid(arrays: NDArray, indexing: str = "xy") -> Sequence[NDArray]:
         * if matrix indexing ij, then each returned array has shape (N1, N2, N3, ..., Nn).
         * if Cartesian indexing xy, then each returned array has shape (N2, N1, N3, ..., Nn).
     """
-    raise NotImplementedError("Working on meshgrid")
+    out = ()
+    shape = np.ones(len(arrays))
+    first_arr = arrays[0]
+    myarrs = ()
+    if indexing == "xy" and len(shape) > 1:
+        # switch 0th and 1st shapes around
+        def mygen(i):
+            if i not in (0, 1):
+                return (j for j in range(len(arrays)) if j != i)
+            else:
+                return (j for j in range(len(arrays)) if j != builtins.abs(i - 1))
+    else:
+        mygen = lambda i: (j for j in range(len(arrays)) if j != i)  # noqa : E731
+
+    for i, a in enumerate(arrays):
+        if len(a.shape) != 1 or a.dtype != first_arr.dtype:
+            raise ValueError("All arrays must be 1D and of same dtype.")
+        shape[i] = a.shape[0]
+        myarrs += (blosc2.expand_dims(a, tuple(mygen(i))),)  # cheap, creates a view
+
+    # handle Cartesian indexing
+    shape = tuple(shape)
+    if indexing == "xy" and len(shape) > 1:
+        shape = (shape[1], shape[0]) + shape[2:]
+
+    # do broadcast
+    for a in myarrs:
+        out += (broadcast_to(a, shape),)
+    return out

tests/ndarray/test_ndarray.py

@@ -409,3 +409,65 @@ def fancy_strided_output(inputs, output_indices, stride=1):
     output_indices = [800, 74, 671, 132, 818]
     out = fancy_strided_output(arr, output_indices, stride=16)
     assert out.shape == (2, 12, 5, 10, 8, 3)
+
+
+dtypes = [np.int32, np.float32, np.float64, np.uint8]
+
+# Shapes for broadcast_to
+broadcast_shapes = [
+    ((10,), (50,), (4,), (3,)),
+    ((8, 6), (16, 12), (4, 3), (1, 3)),
+    ((2, 6), (2, 30), (3, 2), (1, 1)),
+    ((1, 1, 3), (2, 4, 3), (1, 1, 2), (1, 1, 1)),
+]
+
+meshgrid_shapes = [
+    ((10, 20), (3,), (1,)),
+    ((8, 6), (4,), (3,)),
+    ((2, 30), (2,), (1,)),
+    ((20, 4, 3), (4,), (1,)),
+]
+
+
+@pytest.mark.parametrize("dtype", dtypes)
+@pytest.mark.parametrize(("src_shape", "dst_shape", "chunks", "blocks"), broadcast_shapes)
+def test_broadcast_to(dtype, src_shape, dst_shape, chunks, blocks):
+    arr_np = np.arange(np.prod(src_shape), dtype=dtype).reshape(src_shape)
+    arr_b2 = blosc2.asarray(arr_np, chunks=chunks, blocks=blocks)
+
+    try:
+        np_broadcast = np.broadcast_to(arr_np, dst_shape)
+        np_error = None
+    except ValueError as e:
+        np_broadcast = None
+        np_error = e
+
+    if np_error is not None:
+        with pytest.raises(type(np_error)):
+            blosc2.broadcast_to(arr_b2, dst_shape)
+    else:
+        b2_broadcast = blosc2.broadcast_to(arr_b2, dst_shape)
+        assert np.array_equal(b2_broadcast[:], np_broadcast)
+
+
+@pytest.mark.parametrize("dtype", dtypes)
+@pytest.mark.parametrize(("shapes", "chunks", "blocks"), meshgrid_shapes)
+@pytest.mark.parametrize("indexing", ["xy", "ij"])
+def test_meshgrid(dtype, shapes, chunks, blocks, indexing):
+    arrays_np = [np.arange(np.prod(shape), dtype=dtype).reshape(shape) for shape in shapes]
+    arrays_b2 = [blosc2.asarray(a, chunks=chunks, blocks=blocks) for a in arrays_np]
+    try:
+        np_grids = np.meshgrid(*arrays_np, indexing=indexing)
+        np_error = None
+    except ValueError as e:
+        np_grids = None
+        np_error = e
+
+    if np_error is not None:
+        with pytest.raises(type(np_error)):
+            blosc2.meshgrid(*arrays_b2, indexing=indexing)
+    else:
+        b2_grids = blosc2.meshgrid(*arrays_b2, indexing=indexing)
+        assert len(b2_grids) == len(np_grids)
+        for g_b2, g_np in zip(b2_grids, np_grids, strict=False):
+            assert np.array_equal(g_b2[:], g_np)