Merge pull request #408 from Blosc/fixLinspace

Fix linspace for incompatible num/shape

Author: FrancescAlted

src/blosc2/ndarray.py

@@ -3266,7 +3266,9 @@ def arange_fill(inputs, output, offset):
 
 
 # Define a numpy linspace-like function
-def linspace(start, stop, num=50, endpoint=True, dtype=np.float64, shape=None, c_order=True, **kwargs: Any):
+def linspace(
+    start, stop, num=None, endpoint=True, dtype=np.float64, shape=None, c_order=True, **kwargs: Any
+):
     """Return evenly spaced numbers over a specified interval.
 
     This is similar to `numpy.linspace` but it returns a `NDArray`
@@ -3308,8 +3310,17 @@ def linspace_fill(inputs, output, offset):
         stop_ = start_ + lout / num * (stop - start)
         output[:] = np.linspace(start_, stop_, lout, endpoint=False, dtype=output.dtype)
 
-    if not shape:
-        shape = (num,)
+    if shape is None or num is None:
+        if shape is None and num is None:
+            raise ValueError("Either `shape` or `num` must be specified.")
+        if shape is None:  # num is not None
+            shape = (num,)
+        else:  # num is none
+            num = math.prod(shape)
+
+    # check compatibility of shape and num
+    if math.prod(shape) != num:
+        raise ValueError("The specified shape is not consistent with the specified num value")
     dtype = _check_dtype(dtype)
 
     if is_inside_new_expr():

tests/ndarray/test_evaluate.py

@@ -27,7 +27,7 @@ def sample_data(request):
     # The jit decorator can work with any numpy or NDArray params in functions
     a = blosc2.linspace(0, 1, shape[0] * shape[1], dtype=dtype, shape=shape)
     b = np.linspace(1, 2, shape[0] * shape[1], dtype=dtype).reshape(shape)
-    c = blosc2.linspace(-10, 10, cshape[0], dtype=dtype, shape=cshape)
+    c = blosc2.linspace(-10, 10, np.prod(cshape), dtype=dtype, shape=cshape)
     return a, b, c, shape
 
 

tests/ndarray/test_jit.py

@@ -26,7 +26,7 @@ def sample_data(request):
     # The jit decorator can work with any numpy or NDArray params in functions
     a = blosc2.linspace(0, 1, shape[0] * shape[1], dtype=dtype, shape=shape)
     b = np.linspace(1, 2, shape[0] * shape[1], dtype=dtype).reshape(shape)
-    c = blosc2.linspace(-10, 10, cshape[0], dtype=dtype, shape=cshape)
+    c = blosc2.linspace(-10, 10, np.prod(cshape), dtype=dtype, shape=cshape)
     return a, b, c, shape
 
 

tests/ndarray/test_ndarray.py

@@ -187,6 +187,14 @@ def test_linspace(ss, shape, dtype, chunks, blocks, endpoint, c_order):
     else:
         # This is chunk order, so testing is more laborious, and not really necessary
         pass
+    with pytest.raises(ValueError):
+        a = blosc2.linspace(start, stop, 10, shape=(20,))  # num incompatible with shape
+    with pytest.raises(ValueError):
+        a = blosc2.linspace(start, stop)  # num or shape should be specified
+    a = blosc2.linspace(start, stop, shape=(20,))  # should have length 20
+    assert a.shape == (20,)
+    a = blosc2.linspace(start, stop, num=20)  # should have length 20
+    assert a.shape == (20,)
 
 
 @pytest.mark.parametrize(("N", "M"), [(10, None), (10, 20), (20, 10)])