Add test with different shapes and operands

Author: martaiborra

blosc2/blosc2_ext.pyx

@@ -1571,7 +1571,6 @@ cdef int general_numba(blosc2_prefilter_params *params):
     cdef int64_t start_ndim[B2ND_MAX_DIM]
     for i in range(nd):
         start_ndim[i] = chunk_ndim[i] * udata.array.chunkshape[i] + block_ndim[i] * udata.array.blockshape[i]
-        # print("start_ndim[", i, "] = ", start_ndim[i])
 
     padding = False
     blockshape = []
@@ -1606,23 +1605,22 @@ cdef int general_numba(blosc2_prefilter_params *params):
             elif isinstance(obj, np.ndarray):
                 inputs.append(obj[start_ndim[0] : start_ndim[0] + blockshape[0]])
             elif isinstance(obj, (int, float, bool, complex)):
-                inputs.append(np.full(blockshape[0], obj, dtype=dtype))
+                inputs.append(obj)
             else:
                 raise ValueError("Unsupported operand")
     else:
         # Get tuple of slices
         l = []
         for i in range(nd):
             l.append(slice(start_ndim[i], start_ndim[i] + blockshape[i]))
-            # print("slice dim ", i, " = inici = ", start_ndim[i], " final = ",  start_ndim[i] + blockshape[i])
         slices = tuple(l)
         for obj, dtype in inputs_tuple:
             if isinstance(obj, blosc2.NDArray):
                 inputs.append(obj[slices])
             elif isinstance(obj, np.ndarray):
                 inputs.append(obj[slices])
             elif isinstance(obj, (int, float, bool, complex)):
-                inputs.append(np.full(blockshape, obj, dtype=dtype))
+                inputs.append(obj)
             else:
                 raise ValueError("Unsupported operand")
 

tests/ndarray/test_numba_expr.py

@@ -12,9 +12,17 @@
 
 import blosc2
 
+
+@nb.jit(nopython=True, parallel=True)
+def numba1p(inputs_tuple, output, offset):
+    x = inputs_tuple[0]
+    output[:] = x + 1
+
+
 @pytest.mark.parametrize(
     "shape, chunks, blocks",
     [
+        # Test different shapes with and without padding
         ((10, 10), (10, 10), (10, 10),),
         ((20, 20), (10, 10), (10, 10),),
         ((20, 20), (10, 10), (5, 5),),
@@ -24,20 +32,74 @@
         ((13, 13), (10, 10), (4, 4),),
     ],
 )
-def test_numba_expr(shape, chunks, blocks):  #, dtype, urlpath, contiguous
-    @nb.jit(nopython=True, parallel=True)
-    def func_numba(inputs_tuple, output, offset):
-        x = inputs_tuple[0]
-        output[:] = x + 1
-
-    # Create a NDArray from a NumPy array
+def test_numba1p(shape, chunks, blocks):
     npa = np.linspace(0, 1, np.prod(shape)).reshape(shape)
     npc = npa + 1
 
-    # a = blosc2.asarray(npa)
-    expr = blosc2.expr_from_udf(func_numba, ((npa, npa.dtype), ), npa.dtype, chunks=chunks, blocks=blocks)
+    expr = blosc2.expr_from_udf(numba1p, ((npa, npa.dtype), ), npa.dtype, chunks=chunks, blocks=blocks)
     res = expr.eval()
+
     tol = 1e-5 if res.dtype is np.float32 else 1e-14
     if res.dtype in (np.float32, np.float64):
         np.testing.assert_allclose(res[...], npc, rtol=tol, atol=tol)
 
+
+@nb.jit(nopython=True, parallel=True)
+def numba2p(inputs_tuple, output, offset):
+    x = inputs_tuple[0]
+    y = inputs_tuple[1]
+    for i in nb.prange(x.shape[0]):
+        for j in nb.prange(x.shape[1]):
+            output[i, j] = x[i, j] ** 2 + y[i, j] ** 2 + 2 * x[i, j] * y[i, j] + 1
+
+
+@pytest.mark.parametrize(
+    "shape, chunks, blocks",
+    [
+        ((20, 20), (10, 10), (5, 5),),
+        ((13, 13, 10), (10, 10, 5), (5, 5, 3),),
+        ((13, 13), (10, 10), (5, 5),),
+    ],
+)
+def test_numba2p(shape, chunks, blocks):
+    npa = np.arange(0, np.prod(shape)).reshape(shape)
+    npb = np.arange(1, np.prod(shape) + 1).reshape(shape)
+    npc = npa**2 + npb**2 + 2 * npa * npb + 1
+
+    b = blosc2.asarray(npb)
+    expr = blosc2.expr_from_udf(numba2p, ((npa, npa.dtype), (b, b.dtype)), npa.dtype, chunks=chunks, blocks=blocks)
+    res = expr.eval()
+
+    np.testing.assert_allclose(res[...], npc)
+
+
+@nb.jit(nopython=True, parallel=True)
+def numba1dim(inputs_tuple, output, offset):
+    x = inputs_tuple[0]
+    y = inputs_tuple[1]
+    z = inputs_tuple[2]
+    output[:] = x + y + z
+
+
+# Test with np.ndarray, blosc2.SChunk and python scalar operands
+@pytest.mark.parametrize(
+    "shape, chunks, blocks",
+    [
+        ((20, ), (10, ), (5, ),),
+        ((23, ), (10, ), (3, ),),
+    ],
+)
+def test_numba1dim(shape, chunks, blocks):
+    npa = np.arange(start=0, stop=np.prod(shape)).reshape(shape)
+    npb = np.arange(start=1, stop=np.prod(shape) + 1).reshape(shape)
+    py_scalar = np.e
+    npc = npa + npb + py_scalar
+
+    b = blosc2.SChunk(data=npb)
+    expr = blosc2.expr_from_udf(numba1dim, ((npa, npa.dtype), (b, npb.dtype), (py_scalar, np.float64)),
+                                np.float64, chunks=chunks, blocks=blocks)
+    res = expr.eval()
+
+    tol = 1e-5 if res.dtype is np.float32 else 1e-14
+    if res.dtype in (np.float32, np.float64):
+        np.testing.assert_allclose(res[...], npc, rtol=tol, atol=tol)