diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 2229cf7f2..1487a6119 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -21,6 +21,7 @@ import re import sys import threading +import warnings from abc import ABC, abstractmethod, abstractproperty from collections import OrderedDict from collections.abc import MutableMapping @@ -1766,6 +1767,153 @@ def _miniexpr_integer_atan2(expression, operands): ) +def _literal_value(node): + """The number a literal AST node stands for (signs included), or None.""" + sign = 1 + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub | ast.UAdd): + sign = -1 if isinstance(node.op, ast.USub) else 1 + node = node.operand + if isinstance(node, ast.Constant) and isinstance(node.value, bool | int | float): + return sign * node.value + return None + + +def _scalar_operand_value(value): + """The number a scalar operand stands for, or None for arrays (and anything else).""" + if hasattr(value, "shape"): + if value.shape != (): + return None + value = value[()] + if isinstance(value, np.generic): + value = value.item() + return value if isinstance(value, bool | int | float) else None + + +def _exact_in_dtype(value, dtype) -> bool: + """Whether `dtype`, and the C double miniexpr keeps literals in, hold the number exactly.""" + if dtype.kind not in "biuf" or float(value) != value: + return False + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dtype.type(value).item() == value + except (OverflowError, ValueError): + return False + + +@lru_cache(maxsize=1024) +def _parse_expression(expression: str) -> ast.Expression | None: + """The AST of an expression (None if it is not Python syntax); the callers do not modify it.""" + try: + return ast.parse(expression, mode="eval") + except SyntaxError: + return None + + +@lru_cache(maxsize=1024) +def _miniexpr_inexact_int_literal(expression: str) -> bool: + """Whether the expression holds an integer literal beyond 2**53. + + miniexpr parses numeric literals into C doubles, which round such integers. + """ + tree = _parse_expression(expression) + return tree is not None and any( + isinstance(node, ast.Constant) + and isinstance(node.value, int) + and not isinstance(node.value, bool) + and float(node.value) != node.value + for node in ast.walk(tree) + ) + + +def _miniexpr_where_matches_numpy(cond: str, operands: dict, dtype) -> bool: # noqa: C901 + """Whether miniexpr computes where(cond, x, y) as NumPy does, once the where() is fused in. + + miniexpr types the literals of the fused expression by its output dtype, and computes in it, + while NumPy evaluates the condition in the dtypes of its operands, with Python scalars as weak + scalars (NEP 50). With a float32 `a`, NumPy compares ``a > 0.1`` against float32(0.1), but + miniexpr compares ``where(a > 0.1, 1, 0)`` against the double 0.1, and adds ``a + b`` in + ``where(a + b == c, 1, 0)`` in a wider type. The two agree when computing in the output dtype + changes nothing: when every operand of the condition has the output dtype, or when the + condition only compares operands that cast safely to it, against literals that their dtypes + hold exactly. Inside where(), miniexpr also passes 64-bit integers through doubles, so + where() on them stays with NumPy. + """ + if dtype is None: + return False + out_dtype = np.dtype(dtype) + dtypes = [out_dtype] + [ + np.dtype(op.dtype) + for op in operands.values() + if hasattr(op, "dtype") and _scalar_operand_value(op) is None + ] + if any(dt.kind in "iu" and dt.itemsize > 4 for dt in dtypes): + return False + tree = _parse_expression(cond) + if tree is None: + return False + + def term(node): + """("literal", number) or ("array", dtype) for a leaf of the condition, else None.""" + value = _literal_value(node) + if value is not None: + return "literal", value + if isinstance(node, ast.Name) and node.id in operands: + op = operands[node.id] + value = _scalar_operand_value(op) + if value is not None: + return "literal", value + if hasattr(op, "dtype"): + return "array", np.dtype(op.dtype) + return None + + comparisons = [] + + def only_compares(node): + """Whether `node` combines comparisons of leaves with & | ~ only (collected into comparisons).""" + if isinstance(node, ast.Compare): + terms = [term(n) for n in (node.left, *node.comparators)] + if len(node.ops) > 1 or None in terms: + return False + comparisons.append(terms) + return True + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitAnd | ast.BitOr | ast.BitXor): + return only_compares(node.left) and only_compares(node.right) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Invert | ast.Not): + return only_compares(node.operand) + if isinstance(node, ast.BoolOp): + return all(only_compares(value) for value in node.values) + leaf = term(node) + if leaf is not None and leaf[0] == "array" and leaf[1].kind == "b": + comparisons.append([leaf]) + return True + return False + + if only_compares(tree.body): + for terms in comparisons: + dtypes = [dt for kind, dt in terms if kind == "array"] + if not all(np.can_cast(dt, out_dtype, "safe") for dt in dtypes): + return False + literals = [value for kind, value in terms if kind == "literal"] + if not all(_exact_in_dtype(value, dt) for value in literals for dt in (*dtypes, out_dtype)): + return False + return True + + # With arithmetic, the condition is computed as in NumPy only in its own dtype + def same_dtype(node): + leaf = term(node) + if leaf is not None: + kind, value = leaf + return value == out_dtype if kind == "array" else _exact_in_dtype(value, out_dtype) + if isinstance(node, ast.Name): # a function name + return True + if isinstance(node, ast.Constant): # strings and other non-numbers + return False + return all(same_dtype(child) for child in ast.iter_child_nodes(node)) + + return same_dtype(tree.body) + + @blosc2.jit_config.jit_execution def fast_eval( # noqa: C901 expression: str | Callable[[tuple, np.ndarray, tuple[int]], None], @@ -1887,13 +2035,28 @@ def fast_eval( # noqa: C901 expr_string_miniexpr = expr_string operands_miniexpr = operands + # Fuse where(cond, x, y) before specializing, so that scalar x/y get inlined as well; + # fusing afterwards would reference a dropped operand and make miniexpr fail to compile. + where_fused = ( + use_miniexpr + and isinstance(expr_string, str) + and not is_dsl + and where is not None + and len(where) == 2 + ) + if where_fused and not _miniexpr_where_matches_numpy(expr_string, operands, dtype): + use_miniexpr = where_fused = False + if where_fused: + expr_string_miniexpr = f"where({expr_string}, _where_x, _where_y)" if use_miniexpr and isinstance(expr_string, str): - try: - expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs(expr_string, operands) - except Exception: - # If specialization fails, keep original expression/operands and let normal checks decide. - expr_string_miniexpr = expr_string - operands_miniexpr = operands + # If specialization fails, keep original expression/operands and let normal checks decide. + with contextlib.suppress(Exception): + expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs( + expr_string_miniexpr, operands + ) + if not is_dsl and _miniexpr_inexact_int_literal(expr_string_miniexpr): + # Typed literals and inlined scalars alike + use_miniexpr = False # Check whether we can use miniexpr if use_miniexpr: @@ -2025,7 +2188,7 @@ def _miniexpr_eligible_operand(op): k: (np.ascontiguousarray(v) if isinstance(v, np.ndarray) else v) for k, v in _pref_ops.items() } - if where is not None and len(where) == 2: + if where is not None and len(where) == 2 and not where_fused: _pref_expr = f"where({_pref_expr}, _where_x, _where_y)" # _cb_anchor keeps the contiguous bitmap alive for the whole # update_data loop (the prefilter dereferences it per block). Pass @@ -3774,6 +3937,9 @@ def __init__(self, new_op): # noqa: C901 op = None self.operands = {"o0": value1} self.expression = "o0" if op is None else f"{op}(o0)" + if op is None and isinstance(value1, blosc2.NDArray): + # A bare array (as the NDArray reductions build): no need to infer the dtype + self._dtype = value1.dtype return value2 = ( blosc2.SimpleProxy(value2) @@ -4301,51 +4467,69 @@ def mean( out = blosc2.asarray(out, **kwargs) return out - @expression_method - def std( - self, - axis=None, - dtype=None, - keepdims=False, - ddof=0, - where=None, - fp_accuracy: blosc2.FPAccuracy | None = None, - **kwargs, - ): + def _can_full_var(self, dtype): + """Whether _full_var computes the variance in `dtype`: for float32/float64 operands, where + miniexpr evaluates it without per-element conversions, into their own dtype (the inlined + mean would not make the deviations of float32 operands float64).""" + own = np.dtype(self.dtype) + return own.kind == "f" and own.itemsize in (4, 8) and (dtype is None or np.dtype(dtype) == own) + + def _full_var(self, dtype, keepdims, fp_accuracy, execution_kwargs): + """Variance (ddof=0) over all the elements, as two full reductions over the operands. + + The mean goes into the second pass as a Python scalar, which gets inlined into the + expression; a (1, ..., 1) array would be an operand to broadcast, and keep that reduction + off its fast paths. d * d rather than d ** 2 spares a pow() per element. Only for + float32/float64 into the same dtype (see _can_full_var). + """ + mean_value = np.asarray(self.mean(dtype=dtype, fp_accuracy=fp_accuracy, **execution_kwargs)).item() + # Fold the sign into the operator: miniexpr reductions mis-evaluate a negative literal on + # the right of a binary operator, as in "(o0 - -0.5)" + dev = self + (-mean_value) if math.copysign(1.0, mean_value) < 0 else self - mean_value + out = (dev * dev).mean(dtype=dtype, fp_accuracy=fp_accuracy, **execution_kwargs) + # Evaluators other than miniexpr may compute against the Python scalar in float64 + out = np.dtype(self.dtype if dtype is None else dtype).type(out) + return np.reshape(out, (1,) * self.ndim) if keepdims else out + + def _var(self, axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, finish=None): + """The variance, as var() returns it; std() passes np.sqrt as `finish`.""" execution_kwargs = blosc2.jit_config.pop_execution_options(kwargs) where = self._normalize_where(where) item = kwargs.pop("item", ()) - if item == (): # fast path - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - expr = (self - mean_value) ** 2 + if item == () and axis is None and where is None and self._can_full_var(dtype): + out = self._full_var(dtype, keepdims, fp_accuracy, execution_kwargs) else: - mean_value = self.mean( + if item == (): # fast path + mean_value = self.mean( + axis=axis, + dtype=dtype, + keepdims=True, + where=where, + fp_accuracy=fp_accuracy, + **execution_kwargs, + ) + expr = (self - mean_value) ** 2 + else: + mean_value = self.mean( + axis=axis, + dtype=dtype, + keepdims=True, + where=where, + item=item, + fp_accuracy=fp_accuracy, + **execution_kwargs, + ) + # TODO: Not optimal because we load the whole slice in memory. Would have to write + # a bespoke var function that executed within slice_eval to avoid this probably. + expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 + out = expr.mean( axis=axis, dtype=dtype, - keepdims=True, + keepdims=keepdims, where=where, - item=item, fp_accuracy=fp_accuracy, **execution_kwargs, ) - # TODO: Not optimal because we load the whole slice in memory. Would have to write - # a bespoke std function that executed within slice_eval to avoid this probably. - expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 - out = expr.mean( - axis=axis, - dtype=dtype, - keepdims=keepdims, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) if ddof != 0: num_elements = ( self.get_num_elements(axis, item, fp_accuracy=fp_accuracy, **execution_kwargs) @@ -4354,9 +4538,11 @@ def std( axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs ) ) - out = np.sqrt(out * num_elements / (num_elements - ddof)) - else: - out = np.sqrt(out) + # Keep the dtype of the variance: with NumPy 1.x, a float32 scalar times a Python int + # becomes float64, and so does a float32 array times an int above 65535 + out = (out * num_elements / (num_elements - ddof)).astype(np.asarray(out).dtype, copy=False) + if finish is not None: + out = finish(out) out2 = kwargs.pop("out", None) if out2 is not None: out2[:] = out @@ -4365,6 +4551,19 @@ def std( out = blosc2.asarray(out, **kwargs) return out + @expression_method + def std( + self, + axis=None, + dtype=None, + keepdims=False, + ddof=0, + where=None, + fp_accuracy: blosc2.FPAccuracy | None = None, + **kwargs, + ): + return self._var(axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, np.sqrt) + @expression_method def var( self, @@ -4376,56 +4575,7 @@ def var( fp_accuracy: blosc2.FPAccuracy | None = None, **kwargs, ): - execution_kwargs = blosc2.jit_config.pop_execution_options(kwargs) - where = self._normalize_where(where) - item = kwargs.pop("item", ()) - if item == (): # fast path - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - expr = (self - mean_value) ** 2 - else: - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - item=item, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - # TODO: Not optimal because we load the whole slice in memory. Would have to write - # a bespoke var function that executed within slice_eval to avoid this probably. - expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 - out = expr.mean( - axis=axis, - dtype=dtype, - keepdims=keepdims, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - if ddof != 0: - num_elements = ( - self.get_num_elements(axis, item, fp_accuracy=fp_accuracy, **execution_kwargs) - if where is None - else where.where(blosc2.ones(self.shape, dtype=np.int64), 0).sum( - axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs - ) - ) - out = out * num_elements / (num_elements - ddof) - out2 = kwargs.pop("out", None) - if out2 is not None: - out2[:] = out - return out2 - if kwargs != {} and not np.isscalar(out): - out = blosc2.asarray(out, **kwargs) - return out + return self._var(axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs) @expression_method def min( diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 4be24af5d..6837a7b46 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -5,7 +5,9 @@ # SPDX-License-Identifier: BSD-3-Clause ####################################################################### +import itertools import math +import operator import pathlib import warnings @@ -2436,3 +2438,218 @@ def test_cumulative_block_layout(monkeypatch, blocks): expected = np.cumsum(data, axis=0) np.testing.assert_array_equal(result[:], expected) np.testing.assert_array_equal(out[:], expected) + + +def _record_miniexpr(monkeypatch): + """Force miniexpr on, and return the list the expressions it compiles get appended to.""" + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + original_set_pref_expr = blosc2.NDArray._set_pref_expr + compiled = [] + + def wrapped_set_pref_expr(self, expression, inputs, *args, **kwargs): + res = original_set_pref_expr(self, expression, inputs, *args, **kwargs) + compiled.append(expression) + return res + + monkeypatch.setattr(blosc2.NDArray, "_set_pref_expr", wrapped_set_pref_expr) + return compiled + + +@pytest.mark.parametrize( + ("dtype", "x", "y"), + [ + (np.float32, "a", 0), + (np.float32, 1.5, "a"), + (np.int32, "a", 0.5), + (np.int16, "a", -1), + (np.float64, "a", "b"), + ], +) +@pytest.mark.parametrize("getitem", [False, True]) +def test_where_scalar_uses_miniexpr(monkeypatch, dtype, x, y, getitem): + compiled = _record_miniexpr(monkeypatch) + na = np.arange(40 * 30, dtype=dtype).reshape(40, 30) + nb = (na // 3).astype(dtype) + a = blosc2.asarray(na, chunks=(16, 30), blocks=(4, 30)) + b = blosc2.asarray(nb, chunks=(16, 30), blocks=(4, 30)) + ops, nops = {"a": a, "b": b}, {"a": na, "b": nb} + expr = blosc2.where(a > 500, ops.get(x, x), ops.get(y, y)) + expected = np.where(na > 500, nops.get(x, x), nops.get(y, y)) + res = expr[:] if getitem else expr.compute()[:] + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + # The scalar is inlined, so miniexpr (rather than the numexpr fallback) did the work + assert len(compiled) == 1 + assert "_where_x" in compiled[0] or "_where_y" in compiled[0] + + +@pytest.mark.parametrize( + "case", ["mask", "mask_and_threshold", "labels", "labels_arrays", "wider_output", "int16"] +) +@pytest.mark.parametrize("getitem", [False, True]) +def test_where_mixed_dtypes_use_miniexpr(monkeypatch, case, getitem): + # Conditions whose operands cast safely to the output dtype stay on miniexpr + compiled = _record_miniexpr(monkeypatch) + rng = np.random.default_rng(1) + na = rng.normal(100, 50, (40, 30)).astype(np.float32) + nb = rng.normal(size=(40, 30)).astype(np.float32) + nm = na > 120 + nl = rng.integers(0, 5, (40, 30)).astype(np.uint8) + ni = rng.integers(-100, 100, (40, 30)).astype(np.int16) + a, b, m, lbl, i = (blosc2.asarray(v, chunks=(16, 30), blocks=(4, 30)) for v in (na, nb, nm, nl, ni)) + cond, x, y, ncond, nx, ny = { + "mask": (m, a, 0, nm, na, 0), + "mask_and_threshold": (m & (a > 100), a, 0, nm & (na > 100), na, 0), + "labels": (lbl == 3, a, 0, nl == 3, na, 0), + "labels_arrays": (lbl == 3, a, b, nl == 3, na, nb), + "wider_output": (a > 100, 1.0, 0.0, na > 100, 1.0, 0.0), + "int16": (i > 5, a, i, ni > 5, na, ni), + }[case] + expr = blosc2.where(cond, x, y) + expected = np.where(ncond, nx, ny) + res = expr[:] if getitem else expr.compute()[:] + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + assert len(compiled) == 1 + assert "where(" in compiled[0] + + +def test_where_float32_condition_inexact_literal(): + # NumPy compares float32 against float32(0.1), so nothing is > 0.1 here. A fused where() must + # not compare against the double 0.1 (or against 0, the literal typed by an int64 output). + a = blosc2.asarray(np.full(12, 0.1, dtype=np.float32), chunks=(6,), blocks=(3,)) + b = blosc2.asarray(np.full(12, 0.5, dtype=np.float32), chunks=(6,), blocks=(3,)) + for x, y in ((1, 0), (a, 0), (a, b), (1.0, b)): + expr = blosc2.where(a > 0.1, x, y) + expected = np.full(12, y if np.isscalar(y) else 0.5) + np.testing.assert_array_equal(expr[:], expected) + np.testing.assert_array_equal(expr.compute()[:], expected) + expr = blosc2.where(a <= 0.1, x, y) + expected = np.full(12, x if np.isscalar(x) else 0.1, dtype=expr.dtype) + np.testing.assert_array_equal(expr[:], expected) + np.testing.assert_array_equal(expr.compute()[:], expected) + + +def test_where_large_uint64_scalar(): + a = blosc2.asarray(np.arange(12, dtype=np.uint64), chunks=(6,), blocks=(3,)) + big = np.uint64(2**63 + 1) + expr = blosc2.where(a > 5, big, a) + expected = np.array(list(range(6)) + [2**63 + 1] * 6, dtype=np.uint64) + for res in (expr[:], expr.compute()[:]): + assert res.dtype == np.uint64 + np.testing.assert_array_equal(res, expected) + + +@pytest.mark.parametrize("dtype", [np.int64, np.uint64]) +@pytest.mark.parametrize( + "case", ["eq_literal", "add_literal", "add_scalar", "where_scalar", "where_literal_condition"] +) +def test_large_integer_constants_exact(dtype, case): + # miniexpr parses literals into C doubles, so integers beyond 2**53 must not reach it + na = (2**53 + np.arange(-6, 6, dtype=np.int64)).astype(dtype) # the default int is int32 on wasm32 + a = blosc2.asarray(na, chunks=(6,), blocks=(3,)) + big = 2**53 + 1 + huge = dtype(2**63 + 1 if dtype == np.uint64 else 2**62 + 1) + expr, expected = { + "eq_literal": (a == big, na == big), + "add_literal": (a + big, na + big), + "add_scalar": (a + huge, na + huge), + "where_scalar": (blosc2.where(a > 2**53, huge, a), np.where(na > 2**53, huge, na)), + "where_literal_condition": (blosc2.where(a == big, a, 0), np.where(na == big, na, 0)), + }[case] + for res in (expr[:], expr.compute()[:]): + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + + +def _values_around(literal, dtype): + """12 values of `dtype` at and next to `literal`, where comparing with it gets decided.""" + dtype = np.dtype(dtype) + if dtype.kind == "f": + v = dtype.type(literal) + values = [v, np.nextafter(v, dtype.type(np.inf)), np.nextafter(v, dtype.type(-np.inf)), 0] + else: + info = np.iinfo(dtype) + lo, hi = math.floor(literal), math.ceil(literal) + values = [v for v in (lo - 1, lo, hi, hi + 1, 0) if info.min <= v <= info.max] + return np.resize(np.array(values, dtype=dtype), 12) + + +def _where_branches(kind, out_dtype): + """NumPy branches x, y for where(): `out_dtype` scalars, arrays, or an array and a Python 0.""" + if kind == "python": + return (1, 0) if out_dtype == "int" else (0.5, 0.0) + dtype = np.dtype(out_dtype) + # 64-bit integers next to 2**53, to check that their values do not go through doubles + start = 2**53 if dtype.kind in "iu" and dtype.itemsize == 8 else 1 + # int64 explicitly: the default integer is int32 on wasm32, where 2**53 would overflow it + values = np.arange(12, dtype=np.int64) + x = (values % 2).astype(dtype) if dtype.kind == "b" else (start + values).astype(dtype) + y = np.zeros(12, dtype=dtype) + return {"scalars": (x[1], y[0]), "arrays": (x, y), "mixed": (x, 0)}[kind] + + +@pytest.mark.skipif(np.__version__.startswith("1."), reason="NumPy < 2.0 has different casting rules") +@pytest.mark.parametrize( + "cond_dtype", [np.float32, np.float64, np.int8, np.int32, np.int64, np.uint8, np.uint64] +) +@pytest.mark.parametrize("branches", ["python", "scalars", "arrays", "mixed"]) +def test_where_literal_condition_matches_numpy(cond_dtype, branches): + # A fused where() gets typed and computed in its output dtype by miniexpr; comparing against + # literals that the condition's dtype (or a double) cannot hold must still give NumPy's results. + literals = [0.1, 0.5, -1, 100, 16777217, 2**31 + 1, 2**53 + 1] + out_dtypes = ( + ["int", "float"] if branches == "python" else ["float32", "float64", "int64", "uint8", "bool"] + ) + failures = [] + for literal, op, out_dtype in itertools.product(literals, (operator.gt, operator.eq), out_dtypes): + nc = _values_around(literal, cond_dtype) + c = blosc2.asarray(nc, chunks=(6,), blocks=(3,)) + nx, ny = _where_branches(branches, out_dtype) + x, y = ( + blosc2.asarray(v, chunks=(6,), blocks=(3,)) if isinstance(v, np.ndarray) else v for v in (nx, ny) + ) + expected = np.where(op(nc, literal), nx, ny) + expr = blosc2.where(op(c, literal), x, y) + for how, res in (("[:]", expr[:]), ("compute()", expr.compute()[:])): + if res.dtype != expected.dtype or not np.array_equal(res, expected): + failures.append(f"{how} {op.__name__}({literal}) -> {out_dtype}: {res} != {expected}") + assert not failures, "\n".join(failures[:10]) + + +@pytest.mark.skipif(np.__version__.startswith("1."), reason="NumPy < 2.0 has different casting rules") +@pytest.mark.parametrize("cond", ["float32_sum", "int8_wraparound"]) +@pytest.mark.parametrize("branches", ["python_int", "python_float", "array_scalar", "arrays"]) +def test_where_arithmetic_condition_matches_numpy(cond, branches): + # miniexpr would compute the condition in the output dtype: float32 sums in float64, int8 sums + # without wrapping around + rng = np.random.default_rng(0) + if cond == "float32_sum": + na = rng.random(24, dtype=np.float32) + nb = rng.random(24, dtype=np.float32) + operands = (na, nb, na + nb) + + def build(a, b, c): + return a + b == c + else: + na = rng.integers(-128, 128, 24).astype(np.int8) + operands = (na,) + + def build(a): + return a + a < 0 + + to_blosc2 = {id(v): blosc2.asarray(v, chunks=(12,), blocks=(4,)) for v in operands} + nx, ny = { + "python_int": (1, 0), + "python_float": (1.0, 0.0), + "array_scalar": (na, 0), + "arrays": (na, operands[-1]), + }[branches] + x, y = (to_blosc2.get(id(v), v) for v in (nx, ny)) + expected = np.where(build(*operands), nx, ny) + expr = blosc2.where(build(*(to_blosc2[id(v)] for v in operands)), x, y) + for res in (expr[:], expr.compute()[:]): + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) diff --git a/tests/ndarray/test_reductions.py b/tests/ndarray/test_reductions.py index df4951068..5c11ef8de 100644 --- a/tests/ndarray/test_reductions.py +++ b/tests/ndarray/test_reductions.py @@ -983,3 +983,49 @@ def test_reductions_where_ndarray_and_lazyexpr(): np.testing.assert_allclose(a.std(axis=0, where=mask), np.std(na, axis=0, where=nmask)) np.testing.assert_allclose(a.min(axis=0, where=mask), np.min(na, axis=0, where=nmask, initial=np.inf)) np.testing.assert_allclose(a.max(axis=0, where=mask), np.max(na, axis=0, where=nmask, initial=-np.inf)) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("op", ["std", "var"]) +@pytest.mark.parametrize("ddof", [0, 1]) +@pytest.mark.parametrize("keepdims", [False, True]) +@pytest.mark.parametrize("offset", [0, 1e4, -1e4]) +@pytest.mark.parametrize("miniexpr", [True, False]) +def test_std_var_full_reduction(monkeypatch, dtype, op, ddof, keepdims, offset, miniexpr): + # Full std/var reduce (x - mean) * (x - mean) with the mean inlined as a scalar + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", miniexpr) + rng = np.random.default_rng(42) + na = (offset + rng.normal(scale=3, size=(37, 40, 50))).astype(dtype) + a = blosc2.asarray(na, chunks=(10, 20, 50), blocks=(5, 10, 25)) + res = getattr(a, op)(ddof=ddof, keepdims=keepdims) + expected = getattr(na.astype(np.float64), op)(ddof=ddof, keepdims=keepdims) + assert np.asarray(res).dtype == getattr(na, op)().dtype + assert np.shape(res) == np.shape(expected) + np.testing.assert_allclose(res, expected, rtol=1e-5 if dtype == np.float32 else 1e-12) + # Same through a lazy expression and the module-level functions + res = getattr(blosc2, op)(a + 1, ddof=ddof, keepdims=keepdims) + assert np.asarray(res).dtype == getattr(na, op)().dtype + np.testing.assert_allclose(res, expected, rtol=1e-5 if dtype == np.float32 else 1e-12) + + +@pytest.mark.parametrize("op", ["std", "var"]) +def test_std_var_full_reduction_wider_dtype(op): + # A float64 result of float32 data takes float64 deviations from the mean, as NumPy does + rng = np.random.default_rng(1) + na = (rng.random(400_000) + 1e6).astype(np.float32) + a = blosc2.asarray(na) + res = getattr(a, op)(dtype=np.float64) + assert np.asarray(res).dtype == np.float64 + np.testing.assert_allclose(res, getattr(np, op)(na, dtype=np.float64), rtol=1e-6) + + +@pytest.mark.parametrize("value", [np.nan, np.inf]) +def test_std_var_full_reduction_nonfinite(value): + na = np.arange(1000, dtype=np.float64).reshape(10, 100) + na[3, 7] = value + a = blosc2.asarray(na, chunks=(5, 50), blocks=(5, 10)) + with np.errstate(invalid="ignore"): + np.testing.assert_array_equal(a.std(), np.std(na)) + np.testing.assert_array_equal(a.var(), np.var(na))