Fix numerical stability of linspace, add __bool__ for LazyExpr

Author: lshaw8317

src/blosc2/__init__.py

@@ -199,6 +199,22 @@ class Tuner(Enum):
     uint64,
 )
 
+DEFAULT_COMPLEX = complex128
+"""
+Default complex floating dtype."""
+
+DEFAULT_FLOAT = float64
+"""
+Default real floating dtype."""
+
+DEFAULT_INT = int64
+"""
+Default integer dtype."""
+
+DEFAULT_INDEX = int64
+"""
+Default indexing dtype."""
+
 
 class Info:
     def __init__(self, **kwargs):
@@ -218,10 +234,10 @@ def __array_namespace_info__() -> Info:
         },
         default_device=None,
         default_dtypes={
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": DEFAULT_FLOAT,
+            "complex floating": DEFAULT_COMPLEX,
+            "integral": DEFAULT_INT,
+            "indexing": DEFAULT_INDEX,
         },
         dtypes={
             "bool": bool_,
@@ -455,6 +471,11 @@ def __array_namespace_info__() -> Info:
     "MIN_HEADER_LENGTH",
     "VERSION_DATE",
     "VERSION_STRING",
+    # Default dtypes
+    "DEFAULT_COMPLEX",
+    "DEFAULT_FLOAT",
+    "DEFAULT_INDEX",
+    "DEFAULT_INT",
     # Mathematical constants
     "e",
     "pi",

src/blosc2/lazyexpr.py

@@ -489,6 +489,11 @@ def to_cframe(self) -> bytes:
         """
         return self.compute().to_cframe()
 
+    def __bool__(self) -> bool:
+        if math.prod(self.shape) != 1:
+            raise ValueError(f"The truth value of a LazyArray of shape {self.shape} is ambiguous.")
+        return bool(self[()])
+
 
 def convert_inputs(inputs):
     if not inputs or len(inputs) == 0:
@@ -2540,7 +2545,7 @@ def __ror__(self, value):
         return self.update_expr(new_op=(value, "|", self))
 
     def __invert__(self):
-        return self.update_expr(new_op=(self, "~", None))
+        return self.update_expr(new_op=(None, "~", self))
 
     def __pow__(self, value):
         return self.update_expr(new_op=(self, "**", value))

src/blosc2/ndarray.py

@@ -3635,7 +3635,7 @@ def ones(shape: int | tuple | list, dtype: np.dtype | str = None, **kwargs: Any)
     dtype('float64')
     """
     if dtype is None:
-        dtype = blosc2.float64
+        dtype = blosc2.DEFAULT_FLOAT
     return full(shape, 1, dtype, **kwargs)
 
 
@@ -3655,11 +3655,11 @@ def arange(
 
     Parameters
     ----------
-    start: int, float, complex or np.number
+    start: int, float
         The starting value of the sequence.
-    stop: int, float, complex or np.number
+    stop: int, float
         The end value of the sequence.
-    step: int, float, complex or np.number
+    step: int, float or None
         Spacing between values.
     dtype: np.dtype or list str
         The data type of the array elements in NumPy format. Default is
@@ -3717,9 +3717,9 @@ def arange_fill(inputs, output, offset):
             raise ValueError("The shape is not consistent with the start, stop and step values")
     if dtype is None:
         dtype = (
-            blosc2.float64
+            blosc2.DEFAULT_FLOAT
             if np.any([np.issubdtype(type(d), float) for d in (start, stop, step)])
-            else blosc2.int64
+            else blosc2.DEFAULT_INT
         )
     dtype = _check_dtype(dtype)
 
@@ -3742,7 +3742,14 @@ def arange_fill(inputs, output, offset):
 
 # Define a numpy linspace-like function
 def linspace(
-    start, stop, num=None, endpoint=True, dtype=np.float64, shape=None, c_order=True, **kwargs: Any
+    start: int | float | complex,
+    stop: int | float | complex,
+    num: int | None = None,
+    dtype=None,
+    endpoint: bool = True,
+    shape=None,
+    c_order: bool = True,
+    **kwargs: Any,
 ) -> NDArray:
     """Return evenly spaced numbers over a specified interval.
 
@@ -3752,24 +3759,28 @@ def linspace(
 
     Parameters
     ----------
-    start: int, float, complex or np.number
+    start: int, float, complex
         The starting value of the sequence.
-    stop: int, float, complex or np.number
+    stop: int, float, complex
         The end value of the sequence.
-    num: int
-        Number of samples to generate.
+    num: int | None
+        Number of samples to generate. Default None.
+    dtype: np.dtype or list str
+        The data type of the array elements in NumPy format. If None, inferred from
+        start, stop, step. Default is None.
     endpoint: bool
         If True, `stop` is the last sample. Otherwise, it is not included.
-    dtype: np.dtype or list str
-        The data type of the array elements in NumPy format. Default is `np.float64`.
     shape: int, tuple or list
         The shape of the final array. If None, the shape will be guessed from `num`.
     c_order: bool
         Whether to store the array in C order (row-major) or insertion order.
         Insertion order means that values will be stored in the array
         following the order of chunks in the array; this is more memory
         efficient, as it does not require an intermediate copy of the array.
-        Default is C order.
+        Default is True.
+    **kwargs: Any
+        Keyword arguments accepted by the :func:`empty` constructor.
+
 
     Returns
     -------
@@ -3779,33 +3790,48 @@ def linspace(
 
     def linspace_fill(inputs, output, offset):
         lout = len(output)
-        start, stop, num = inputs
+        start, stop, num, endpoint = inputs
+        # if num = 1 do nothing
+        step = (stop - start) / (num - 1) if endpoint and num > 1 else (stop - start) / num
         # Compute proper start and stop values for the current chunk
-        start_ = start + offset[0] / num * (stop - start)
-        stop_ = start_ + lout / num * (stop - start)
-        output[:] = np.linspace(start_, stop_, lout, endpoint=False, dtype=output.dtype)
-
-    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)
+        # except for 0th iter, have already included start_ in prev iter
+        start_ = start + offset[0] * step if offset[0] == 0 else start + (offset[0] + 1) * step
+        stop_ = start_ + lout * step
+        if offset[0] + lout == num:  # reached end
+            output[:] = np.linspace(start_, stop, lout, endpoint=endpoint, dtype=output.dtype)
+        else:  # always include start and stop
+            output[:] = np.linspace(start_, stop_, lout, endpoint=True, dtype=output.dtype)
+
+    if num < 0:
+        raise ValueError("num must be nonnegative.")
+
+    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")
+
+    if dtype is None:
+        dtype = (
+            blosc2.DEFAULT_COMPLEX
+            if np.any([np.issubdtype(type(d), complex) for d in (start, stop)])
+            else blosc2.DEFAULT_FLOAT
+        )
+
     dtype = _check_dtype(dtype)
 
-    if is_inside_new_expr():
+    if is_inside_new_expr() or num == 0:
         # We already have the dtype and shape, so return immediately
-        return blosc2.zeros(shape, dtype=dtype)
+        return blosc2.zeros(shape, dtype=dtype)  # will return empty array for num == 0
 
     lshape = (math.prod(shape),)
-    if endpoint:
-        stop += (stop - start) / (num - 1)
-    inputs = (start, stop, num)
+
+    inputs = (start, stop, num, endpoint)
     lazyarr = blosc2.lazyudf(linspace_fill, inputs, dtype=dtype, shape=lshape)
     if len(shape) == 1:
         # C order is guaranteed, and no reshape is needed