Give lazy expressions NumPy's result dtype for division and float functions

src/blosc2/lazyexpr.py

@@ -358,6 +358,23 @@ def _find_constructor_call(expression: str, constructor: str) -> re.Match | None
     "//": operator.floordiv,
     "%": operator.mod,
 }
+# NumPy operations of the numeric operators and functions whose result dtype is not just the
+# promoted operand dtype: true division and the float functions turn integers (and bools) into
+# floats, and //, %, ** and the shifts, which NumPy has no bool loops for, turn bools into int8
+_result_dtype_ufuncs = {
+    "/": operator.truediv,
+    "//": operator.floordiv,
+    "%": operator.mod,
+    "**": operator.pow,  # not np.power: NumPy squares for x ** 2, which keeps bools int8
+    "<<": operator.lshift,
+    ">>": operator.rshift,
+    "arctan2": np.arctan2,
+    "copysign": np.copysign,
+    "hypot": np.hypot,
+    "nextafter": np.nextafter,
+    "pow": np.power,
+    "power": np.power,
+}
 logical_ops = ["&", "|", "^", "~"]
 not_complex_ops = ["maximum", "minimum", "<", "<=", ">", ">="]
 funcs_2args = (
@@ -3624,12 +3641,24 @@ def fuse_expressions(expr, new_base, dup_op):
     return new_expr
 
 
+def _operand_dtype(value):
+    """blosc2.result_type(value), without building a NumPy array for an array operand."""
+    dtype = None if np.isscalar(value) else getattr(value, "dtype", None)
+    if dtype is not None:
+        dtype = convert_dtype(dtype)
+        if dtype.kind != "V":
+            # np.result_type() gives the native byte order, which the dtype checks below compare against
+            return np.result_type(dtype)
+    # Structured and void dtypes too: blosc2.result_type() rejects some of them
+    return blosc2.result_type(value)
+
+
 def check_dtype(op, value1, value2):
     if op in ("contains", "startswith", "endswith"):
         return np.dtype(np.bool_)
 
-    v1_dtype = blosc2.result_type(value1)
-    v2_dtype = v1_dtype if value2 is None else blosc2.result_type(value2)
+    v1_dtype = _operand_dtype(value1)
+    v2_dtype = v1_dtype if value2 is None else _operand_dtype(value2)
     if op in not_complex_ops and (v1_dtype == np.complex128 or v2_dtype == np.complex128):
         # Ensure that throw exception for functions which don't support complex args
         raise ValueError(f"Invalid operand type for {op}: {v1_dtype, v2_dtype}")
@@ -3642,13 +3671,23 @@ def check_dtype(op, value1, value2):
         if v2_dtype not in (np.bool_, np.int32, np.int64):
             raise ValueError(f"Invalid operand type for {op}: {v2_dtype}")
 
-    if op == "/":
-        if v1_dtype == np.int32 and v2_dtype == np.int32:
-            return blosc2.float32
-        if np.issubdtype(v1_dtype, np.integer) and np.issubdtype(v2_dtype, np.integer):
-            return blosc2.float64
+    if op == "/" and v1_dtype == np.int32 and v2_dtype == np.int32:
+        return blosc2.float32
+
+    ufunc = _result_dtype_ufuncs.get(op)
+    if ufunc is not None and value2 is not None and v1_dtype.kind in "biufc" and v2_dtype.kind in "biufc":
+        # Ask NumPy on empty operands, but with a scalar next to an array as it is, so that NumPy
+        # promotes it as usual; what it rejects is left to the evaluation
+        scalars = np.isscalar(value1) and np.isscalar(value2)
+        probes = [
+            v if np.isscalar(v) and not scalars else np.empty(0, dt)
+            for v, dt in ((value1, v1_dtype), (value2, v2_dtype))
+        ]
+        with contextlib.suppress(TypeError, ValueError, OverflowError):
+            return ufunc(*probes).dtype
 
     if op == "arctan2":
+        # Except for arctan2, which rejects what NumPy rejects (complex, datetimes, strings) here
         return np.arctan2(np.empty(0, v1_dtype), np.empty(0, v2_dtype)).dtype
 
     if (v1_dtype.kind in "Mm" or v2_dtype.kind in "Mm") and op in _datetime_binops:

tests/ndarray/test_lazyexpr.py

@@ -45,6 +45,125 @@ def test_arctan2_preserves_fractional_results(dtype, fused):
     np.testing.assert_allclose(blosc2.lazyexpr("arctan2(a, b)")[:], expected, rtol=1e-6)
 
 
+_RESULT_DTYPE_EXPRESSIONS = [
+    "hypot({n}, b)",
+    "copysign({n}, b)",
+    "nextafter({n}, b)",
+    "arctan2({n}, b)",
+    "hypot(a, b * 1)",
+    "hypot({n}, 2)",
+    "arctan2({n}, 2.5)",
+    "(a < b) / b",
+    "a / 2",
+    "a / b",
+    "{n} / b",
+    "{n} // b",
+    "{n} % b",
+    "a ** b",
+    "{n} ** 2",
+    "sin({n})",
+    "sqrt(a)",
+]
+
+
+# Known wrong values on WebAssembly: there compute() and slices also run the bundled miniexpr, which
+# computes these in integer arithmetic (0.0 for NumPy's 0.5). strict, so that a fix shows up as XPASS
+_MINIEXPR_TRUE_DIVISION = (
+    "miniexpr computes integer true division in integer arithmetic: "
+    "https://github.com/Blosc/miniexpr/issues/26"
+)
+_MINIEXPR_FLOAT_FUNCTIONS = (
+    "miniexpr computes hypot() and nextafter() of integers in integer arithmetic: "
+    "https://github.com/Blosc/miniexpr/issues/27"
+)
+_WASM_KNOWN_FAILURES = {
+    "hypot({n}, b)": ((np.uint16, np.int32, np.int64), _MINIEXPR_FLOAT_FUNCTIONS),
+    "nextafter({n}, b)": ((np.int32, np.int64), _MINIEXPR_FLOAT_FUNCTIONS),
+    "(a < b) / b": ((np.int8, np.uint16, np.int32, np.int64), _MINIEXPR_TRUE_DIVISION),
+    "a / b": ((np.int8, np.uint16, np.int64), _MINIEXPR_TRUE_DIVISION),
+    "{n} / b": ((np.int8, np.uint16, np.int64), _MINIEXPR_TRUE_DIVISION),
+}
+
+
+@pytest.mark.parametrize("dtype", [np.bool_, np.int8, np.uint16, np.int32, np.int64])
+@pytest.mark.parametrize("expression", _RESULT_DTYPE_EXPRESSIONS)
+@pytest.mark.parametrize("build", ["functions", "string"])
+def test_lazyexpr_result_dtype_matches_numpy(dtype, expression, build, request):
+    # Float functions and true division of integer (or bool) subexpressions were given an integer
+    # (or bool) result dtype, and bool //, % and ** a bool one, so their results were truncated
+    if dtype is np.int32 and expression in ("a / b", "{n} / b"):
+        pytest.skip("int32 / int32 gives float32 on purpose")
+    if "nextafter" in expression and dtype not in (np.int32, np.int64):
+        pytest.skip("a float16 or float32 nextafter() is computed in float64 on every path")
+    known_dtypes, reason = _WASM_KNOWN_FAILURES.get(expression, ((), None))
+    if blosc2.IS_WASM and dtype in known_dtypes:
+        request.applymarker(pytest.mark.xfail(reason=reason, strict=True))
+    # A subexpression of the operand dtype (bool + bool is an or in NumPy)
+    expression = expression.format(n="(a & b)" if dtype is np.bool_ else "(a + b)")
+    na = (np.arange(1200).reshape(30, 40) % 5).astype(dtype)
+    nb = (np.arange(1200).reshape(30, 40) % 3 + 1).astype(dtype)
+    a = blosc2.asarray(na, chunks=(10, 40), blocks=(5, 40))
+    b = blosc2.asarray(nb, chunks=(10, 40), blocks=(5, 40))
+    namespace = {"hypot": np.hypot, "copysign": np.copysign, "nextafter": np.nextafter}
+    namespace.update(arctan2=np.arctan2, sin=np.sin, sqrt=np.sqrt)
+    with np.errstate(all="ignore"):
+        expected = eval(expression, namespace, {"a": na, "b": nb})
+        if dtype is np.int32 and expression == "a / 2" and blosc2.result_type(2) == np.int32:
+            # A Python int is int32 here (as on wasm32), so this is int32 / int32, which gives float32
+            # on purpose; elsewhere it is int32 / int64, float64 as in NumPy
+            expected = expected.astype(np.float32)
+        if build == "string":
+            expr = blosc2.lazyexpr(expression, {"a": a, "b": b})
+        else:
+            expr = eval(expression, {name: getattr(blosc2, name) for name in namespace}, {"a": a, "b": b})
+        assert expr.dtype == expected.dtype
+        for res in (expr[:], expr[...], expr.compute()[:], expr[:7]):
+            assert res.dtype == expected.dtype
+        # Values on the paths whose typing follows NumPy here; expr[:] and reductions of integer
+        # true division and two-argument float functions still compute in integer arithmetic in
+        # the bundled miniexpr
+        rtol = 1e-3 if expected.dtype == np.float16 else 1e-6
+        np.testing.assert_allclose(expr.compute()[:], expected, rtol=rtol)
+        np.testing.assert_allclose(expr[:7], expected[:7], rtol=rtol)
+
+
+@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float64])
+@pytest.mark.parametrize("op", ["&", "|", "^", "<<", ">>", "+", "*", "/", "//", "%", "**", "<", "hypot"])
+def test_lazyexpr_non_native_byte_order(dtype, op):
+    # The operand dtype checks compare against native dtypes: a non-native operand must give the
+    # same result dtype as a native one (no "Invalid operand type for &: >i4", and >i4 / >i4
+    # keeps the float32 of int32 / int32)
+    if dtype is np.float64 and op in ("&", "|", "^", "<<", ">>"):
+        pytest.skip("no float bitwise operators")
+    native = np.dtype(dtype)
+    swapped = native.newbyteorder()
+    na = np.arange(1, 7, dtype=native)
+    nb = np.array([1, 2, 1, 3, 1, 2], dtype=native)
+
+    def build(a, b):
+        return blosc2.hypot(a, b) if op == "hypot" else eval(f"a {op} b")
+
+    reference = build(blosc2.asarray(na), nb)
+    if op == "/" and dtype is np.int32:
+        assert np.dtype(reference.dtype) == np.float32
+    expr = build(blosc2.asarray(na), nb.astype(swapped))
+    assert expr.dtype == reference.dtype
+    np.testing.assert_array_equal(expr[:], reference[:])
+    np.testing.assert_array_equal(expr.compute()[:], reference.compute()[:])
+    # Only the dtype for non-native NDArray operands: their values are a separate problem
+    sa, sb = na.astype(swapped), nb.astype(swapped)
+    for a, b in ((na, sb), (sa, nb), (sa, sb)):
+        assert build(blosc2.asarray(a), blosc2.asarray(b)).dtype == reference.dtype
+
+
+@pytest.mark.parametrize("other", [np.ones(4, dtype=np.int32), 2, 2.5])
+def test_arctan2_rejects_complex_when_built(other):
+    c = blosc2.asarray(np.ones(4, dtype=np.complex128))
+    for x in (c, c + c):
+        with pytest.raises(TypeError):
+            blosc2.arctan2(x, other)
+
+
 _UNARY_FUNCTIONS = [
     "sin",
     "cos",