diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 2229cf7f2..e4b46f9c8 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 @@ -68,6 +69,7 @@ record_expression_metadata, validate_operands, ) +from .ndarray import blocks_contiguous from .prepared_evaluate import remember_evaluation from .proxy_source import convert_dtype from .utils import ( @@ -318,6 +320,9 @@ def _get_result(expression, chunk_operands, eval_kwargs, where=None, indices=Non functions = blosc2_funcs _TRANSIENT_MASK_CPARAMS = blosc2.CParams(codec=blosc2.Codec.LZ4, clevel=5, filters=[blosc2.Filter.SHUFFLE]) _constructor_call_patterns = {name: re.compile(rf"\b{re.escape(name)}\s*\(") for name in constructors} +_any_constructor_call_pattern = re.compile( + rf"\b(?:{'|'.join(re.escape(name) for name in constructors)})\s*\(" +) def _restore_code_unit_shuffle(cparams: blosc2.CParams, dtype) -> None: @@ -344,6 +349,10 @@ def _has_constructor_call(expression: str, constructor: str) -> bool: return _constructor_call_patterns[constructor].search(expression) is not None +def _has_any_constructor_call(expression: str) -> bool: + return _any_constructor_call_pattern.search(expression) is not None + + def _find_constructor_call(expression: str, constructor: str) -> re.Match | None: return _constructor_call_patterns[constructor].search(expression) @@ -1379,7 +1388,7 @@ def do_slices_intersect(slice1: list | tuple, slice2: list | tuple) -> bool: def get_chunk(arr, info, nchunk): - reduc, aligned, low_mem, chunks_idx = info + reduc, low_mem, chunks_idx = info if low_mem: # We don't want to uncompress the chunk, so keep it compressed and @@ -1403,14 +1412,6 @@ def get_chunk(arr, info, nchunk): slice(c * s, min((c + 1) * s, shape[i])) for i, (c, s) in enumerate(zip(coords, chunks, strict=True)) ) - chunks_ = tuple(s.stop - s.start for s in slice_) - - if aligned: - # Decompress the whole chunk and return it - buff = arr.schunk.decompress_chunk(nchunk) - bsize = arr.dtype.itemsize * math.prod(chunks_) - return np.frombuffer(buff[:bsize], dtype=arr.dtype).reshape(chunks_) - return arr[slice_] @@ -1440,10 +1441,10 @@ def read_chunks_worker(arrs, info, queue, stop): max_workers = max(1, min(len(arrs), int(getattr(blosc2, "nthreads", 1) or 1))) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: my_chunk_iter = range(arrs[0].schunk.nchunks) - if len(info) == 5: + if len(info) == 4: if info[-1] is not None: my_chunk_iter = sliced_chunk_iter(chunks_, (), shape, axis=info[-1], nchunk=True) - info = info[:4] + info = info[:3] for i, nchunk in enumerate(my_chunk_iter): futures = [executor.submit(get_chunk, arr, info, nchunk) for arr in arrs] # result() keeps operand order and propagates the first exception raised @@ -1504,12 +1505,43 @@ def read_nchunk(arrs, info): iter_chunks = None +def _chunk_buffers(operands, shape, chunks): + """Reusable chunk-sized buffers for the NDArray operands of `shape` and `chunks` whose + decompressed chunks are C-order. + + That is the case when blocks are C-contiguous runs of their chunk (as with the default + partitions), so a plain decompress_chunk() gives the (padded) chunk, and is cheaper than a + slice read. + """ + return { + key: np.empty(value.chunks, dtype=value.dtype) + for key, value in operands.items() + if isinstance(value, blosc2.NDArray) + and value.ndim > 0 + and value.shape == shape + and value.chunks == chunks + and 0 not in value.blocks + and blocks_contiguous(value.chunks, value.blocks) + } + + def fill_chunk_operands( - operands, slice_, chunks_, full_chunk, aligned, nchunk, iter_disk, chunk_operands, reduc=False, axis=None + operands, + slice_, + chunks_, + full_chunk, + nchunk, + iter_disk, + chunk_operands, + reduc=False, + axis=None, + chunk_buffers=None, ): """Retrieve the chunk operands for evaluating an expression. This function provides an optimized path for full chunks and a slower path for partial chunks. + Operands with a buffer in `chunk_buffers` (see `_chunk_buffers`) get their chunk decompressed + into it, whether the chunk is full or not. """ global iter_chunks @@ -1521,28 +1553,16 @@ def fill_chunk_operands( if iter_chunks is None: # Initialize the iterator for reading the chunks # Take any operand (all should have the same shape and chunks) - key, arr = next(iter(operands.items())) + arr = next(iter(operands.values())) chunks_idx, _ = get_chunks_idx(arr.shape, arr.chunks) - info = (reduc, aligned[key], low_mem, chunks_idx, axis) + info = (reduc, low_mem, chunks_idx, axis) iter_chunks = read_nchunk(list(operands.values()), info) # Run the asynchronous file reading function from a synchronous context chunks = next(iter_chunks) for i, (key, value) in enumerate(operands.items()): - # Chunks are already decompressed, so we can use them directly - if not low_mem: - if full_chunk: - chunk_operands[key] = chunks[i] - else: - chunk_operands[key] = value[slice_] - continue - # Otherwise, we need to decompress them - if aligned[key]: - buff = blosc2.decompress2(chunks[i]) - bsize = value.dtype.itemsize * math.prod(chunks_) - chunk_operands[key] = np.frombuffer(buff[:bsize], dtype=value.dtype).reshape(chunks_) - else: - chunk_operands[key] = value[slice_] + # Chunks are already decompressed (unless low_mem), so full ones can be used directly + chunk_operands[key] = chunks[i] if full_chunk and not low_mem else value[slice_] return # Get the starts and stops for the slice @@ -1556,6 +1576,11 @@ def fill_chunk_operands( if value.shape == (): chunk_operands[key] = value[()] continue + buf = chunk_buffers.get(key) if chunk_buffers else None + if buf is not None: + value.schunk.decompress_chunk(nchunk, buf) + chunk_operands[key] = buf if full_chunk else buf[tuple(slice(0, n) for n in chunks_)] + continue if not full_chunk or not isinstance(value, blosc2.NDArray): # The chunk is not a full one, or has padding, or is not a blosc2.NDArray, # so we need to go the slow path @@ -1571,13 +1596,7 @@ def fill_chunk_operands( value.get_slice_numpy(chunk_operands[key], (starts, stops)) continue - if aligned[key]: - # Decompress the whole chunk and store it - buff = value.schunk.decompress_chunk(nchunk) - bsize = value.dtype.itemsize * math.prod(chunks_) - chunk_operands[key] = np.frombuffer(buff[:bsize], dtype=value.dtype).reshape(chunks_) - else: - chunk_operands[key] = value[slice_] + chunk_operands[key] = value[slice_] def _apply_jit_backend_pragma(expression: str, inputs: dict, jit_backend: str | None) -> str: @@ -1766,6 +1785,178 @@ 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) + + +def _c_run_axis(outer, inner) -> int: + """Axis along which `inner` tiles are C-contiguous runs of `outer`, or -1 if they are not. + + That is the case when the leading dims of the tile are 1, then one dim is free, and the rest + span the container. + """ + k = 0 + while k < len(inner) - 1 and inner[k] == 1: + k += 1 + return k if tuple(inner[k + 1 :]) == tuple(outer[k + 1 :]) else -1 + + +def _chunks_tile_flat(shape, chunks, blocks) -> bool: + """Whether the uncompressed chunks of an array, in order, are consecutive runs of its C-order buffer. + + Blocks must be C-contiguous runs of their chunk, and chunks C-contiguous runs of the array; only + the very last chunk may be partial, so that its padding spills past the end of the array instead + of into another chunk. + """ + if len(shape) == 0 or 0 in shape or 0 in chunks or 0 in blocks or not blocks_contiguous(chunks, blocks): + return False + k = _c_run_axis(shape, chunks) + return k >= 0 and (shape[k] % chunks[k] == 0 or math.prod(shape[:k]) == 1) + + @blosc2.jit_config.jit_execution def fast_eval( # noqa: C901 expression: str | Callable[[tuple, np.ndarray, tuple[int]], None], @@ -1865,11 +2056,7 @@ def fast_eval( # noqa: C901 blocks = kwargs.pop("blocks", None) if blocks is None: blocks = basearr.blocks - # Check whether the partitions are aligned and behaved - aligned = { - k: False if not hasattr(k, "chunks") else blosc2.are_partitions_aligned(k.shape, k.chunks, k.blocks) - for k in operands - } + # Check whether the partitions are behaved behaved = blosc2.are_partitions_behaved(shape, chunks, blocks) # Check that all operands are NDArray for fast path @@ -1887,13 +2074,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: @@ -1953,14 +2155,18 @@ def _miniexpr_eligible_operand(op): ) numpy_output = ( (getitem or isinstance(out, np.ndarray)) - and len(shape) == 1 - and chunks[0] % blocks[0] == 0 + # Each chunk's payload must be a consecutive run of the C-order result + and _chunks_tile_flat(shape, chunks, blocks) and not prune_chunks and (out is None or not is_dsl) # Rebinding an existing LazyExpr can leave unused operand names. # Retain its established fallback rather than compiling extra inputs. and (is_dsl or set(operands_miniexpr) == set(get_expr_operands(expr_string))) - and (out is None or (out.dtype == dtype and out.flags.writeable)) + and ( + out is None + # A 1-D out may be strided; an N-D one is written through its flat view + or (out.dtype == dtype and out.flags.writeable and (out.ndim == 1 or out.flags.c_contiguous)) + ) ) if not (all_eligible_miniexpr and (out is None or numpy_output)): use_miniexpr = False @@ -2025,7 +2231,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 @@ -2040,6 +2246,8 @@ def _miniexpr_eligible_operand(op): # Data to compress is fetched from operands, so it can be uninitialized here data = np.empty(res_eval.schunk.chunksize, dtype=np.uint8) if numpy_output: + chunk_nitems = math.prod(chunks) + out_flat = None if out is None else (out if out.ndim == 1 else out.reshape(-1)) output_storage = None if out is None: # Reserve a header prefix and at most one chunk of tail @@ -2070,8 +2278,9 @@ def _miniexpr_eligible_operand(op): if chunk_has_candidate is not None and not chunk_has_candidate[nchunk]: continue if numpy_output: - start = nchunk * chunks[0] - count = min(chunks[0], shape[0] - start) + # Chunks are consecutive runs of the C-order result (see _chunks_tile_flat) + start = nchunk * chunk_nitems + count = min(chunk_nitems, out.size - start) if output_storage is not None: byte_start = nchunk * data.size chunk_data = output_storage[ @@ -2090,7 +2299,7 @@ def _miniexpr_eligible_operand(op): values = np.frombuffer( chunk_data, dtype=dtype, count=count, offset=blosc2.EXTENDED_HEADER_LENGTH ) - out[start : start + count] = values + out_flat[start : start + count] = values else: res_eval.schunk.update_data(nchunk, data, copy=False) # _cb_anchor (if any) stays referenced until fast_eval returns, which @@ -2138,6 +2347,7 @@ def _miniexpr_eligible_operand(op): return res_eval chunk_operands = {} + chunk_buffers = None if iter_disk else _chunk_buffers(operands, shape, chunks) # Check which chunks intersect with _slice all_chunks = get_intersecting_chunks((), shape, chunks) # if _slice is (), returns all chunks for nchunk, chunk_slice in enumerate(all_chunks): @@ -2147,7 +2357,15 @@ def _miniexpr_eligible_operand(op): full_chunk = chunks_ == chunks # slice is same as chunk fill_chunk_operands( - operands, cslice, chunks_, full_chunk, aligned, nchunk, iter_disk, chunk_operands + operands, + cslice, + chunks_, + full_chunk, + nchunk, + iter_disk, + chunk_operands, + # Partial chunks keep their contiguous slice reads, which UDFs may rely on + chunk_buffers=chunk_buffers if full_chunk else None, ) # Since evaluate_chunk() can return a dtype larger than the one computed in the expression, @@ -2936,7 +3154,6 @@ def reduce_slices( # noqa: C901 same_chunks = all(operand.chunks == o.chunks for o in operands.values() if hasattr(o, "chunks")) same_blocks = all(operand.blocks == o.blocks for o in operands.values() if hasattr(o, "blocks")) fast_path = same_shape and same_chunks and same_blocks and (0 not in operand.chunks) - aligned = dict.fromkeys(operands.keys(), False) iter_disk = False if fast_path: chunks = operand.chunks @@ -3093,16 +3310,19 @@ def reduce_slices( # noqa: C901 intersecting_chunks = get_intersecting_chunks(_slice, shape, chunks) out_init = False res_out_init = False - ratio = ( - np.ceil(np.asarray(shape) / np.asarray(chunks)).astype(np.int64) - if 0 not in chunks - else np.asarray(shape) + # Flat chunk index = sum(chunk coordinate * stride) over the chunk grid + chunk_strides = [] + if 0 not in chunks: + nchunks_dim, _ = get_chunks_idx(shape, chunks) + chunk_strides = [math.prod(nchunks_dim[i + 1 :]) for i in range(len(nchunks_dim))] + chunk_buffers = ( + _chunk_buffers(operands, shape, chunks) if fast_path and not iter_disk and _slice == () else None ) for chunk_slice in intersecting_chunks: cslice = chunk_slice.raw nchunk = ( - builtins.sum(c.start // chunks[i] * np.prod(ratio[i + 1 :]) for i, c in enumerate(cslice)) + builtins.sum(c.start // chunks[i] * chunk_strides[i] for i, c in enumerate(cslice)) if 0 not in chunks else 0 ) @@ -3112,10 +3332,10 @@ def reduce_slices( # noqa: C901 cslice = step_handler(cslice, _slice) offset = tuple(s.start for s in cslice) # offset for the udf starts = [s.start if s.start is not None else 0 for s in cslice] - unit_steps = np.all([s.step == 1 for s in cslice]) + unit_steps = builtins.all(s.step == 1 for s in cslice) cslice_shape = tuple(s.stop - s.start for s in cslice) - # get local index of part of out that is to be updated - cslice_subidx = ndindex.ndindex(cslice).as_subindex(_slice).raw # if _slice is (), just gives cslice + # get local index of part of out that is to be updated (with no _slice, that is cslice) + cslice_subidx = cslice if _slice == () else ndindex.ndindex(cslice).as_subindex(_slice).raw if _slice == () and fast_path and unit_steps: # Fast path full_chunk = cslice_shape == chunks @@ -3124,12 +3344,12 @@ def reduce_slices( # noqa: C901 cslice, cslice_shape, full_chunk, - aligned, nchunk, iter_disk, chunk_operands, reduc=True, axis=reduce_args["axis"] if np.isscalar(reduce_args["axis"]) else None, + chunk_buffers=chunk_buffers, ) else: get_chunk_operands(operands, cslice, chunk_operands, shape) @@ -3774,6 +3994,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) @@ -4057,7 +4280,7 @@ def shape(self): return None # Operands shape can change, so we always need to recompute this - if any(_has_constructor_call(self.expression, constructor) for constructor in constructors): + if _has_any_constructor_call(self.expression): # might have an expression with pure constructors opshapes = {k: v if not hasattr(v, "shape") else v.shape for k, v in self.operands.items()} _shape = infer_shape(self.expression, opshapes) # infer shape, includes constructors @@ -4301,51 +4524,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 +4595,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 +4608,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 +4632,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( @@ -4831,7 +5038,7 @@ def _compute_expr(self, item, kwargs): # noqa: C901 return chunked_eval(lazy_expr.expression, lazy_expr.operands, item, **kwargs) - if any(_has_constructor_call(self.expression, constructor) for constructor in constructors): + if _has_any_constructor_call(self.expression): expression = self.expression newexpr = expression newops = self.operands.copy() @@ -6205,7 +6412,7 @@ def lazyexpr( operands = seek_operands(operand_set, local_dict, global_dict, _frame_depth=_frame_depth) else: # No operands found in the expression. Maybe a constructor? - constructor = any(_has_constructor_call(expression, constructor) for constructor in constructors) + constructor = _has_any_constructor_call(expression) if not constructor and parse_expression(expression).names: # Approved constant-only numerical expressions are valid too; # parsing, rather than Python eval, decides whether to admit them. diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 5284a0a61..e7210e1d9 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -307,6 +307,21 @@ def are_partitions_aligned(shape, chunks, blocks): return builtins.all(c % b == 0 for c, b in zip(chunks, blocks, strict=True)) +def blocks_contiguous(chunks, blocks): + """Whether the blocks are C-contiguous runs of their chunk (no block may be 0 long). + + That is the case only if their leading dims are 1, then one dim divides the chunk's, and the + rest are equal to the chunk's; same rule as nchunk_fastpath() in C-Blosc2's b2nd.c. + """ + k = 0 + for i, (size, unit) in enumerate(zip(chunks, blocks, strict=True)): + if size % unit != 0 or (i > k and size != unit): + return False + if i == k and unit == 1: + k += 1 + return True + + def are_partitions_behaved(shape, chunks, blocks): """ Check if the partitions defined by chunks and blocks are well-behaved with respect to the shape. @@ -324,18 +339,6 @@ def check_contiguity(container, part): return False return builtins.all(size % unit == 0 for size, unit in zip(container[:-1], part[:-1], strict=True)) - def blocks_contiguous(chunks, blocks): - # A block is a C-contiguous run of its chunk only if its leading dims are 1, - # then one dim divides the chunk's, and the rest are equal to the chunk's. - # Same rule as nchunk_fastpath() in C-Blosc2's b2nd.c. - k = 0 - for i, (size, unit) in enumerate(zip(chunks, blocks, strict=True)): - if size % unit != 0 or (i > k and size != unit): - return False - if i == k and unit == 1: - k += 1 - return True - return blocks_contiguous(chunks, blocks) and check_contiguity(shape, chunks) diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 4be24af5d..32f0bc230 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,270 @@ 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) + + +@pytest.mark.parametrize( + ("shape", "chunks", "blocks", "direct"), + [ + ((100, 100), (100, 100), (25, 100), True), + ((10_000,), (3000,), (1000,), True), # partial last chunk + ((5,), (8,), (4,), True), # a single, partial chunk + ((37, 50, 60), (5, 50, 60), (1, 25, 60), True), + ((4, 12, 30), (1, 4, 30), (1, 2, 30), True), + ((1, 10, 30), (1, 4, 30), (1, 2, 30), True), # partial, but the very last chunk + ((4, 10, 30), (1, 4, 30), (1, 2, 30), False), # partial chunks inside the array + ((40, 30), (16, 30), (4, 10), False), # blocks not contiguous in their chunk + ((40, 30), (16, 15), (4, 15), False), # chunks not contiguous in the array + ((24, 20, 18), (8, 10, 9), (4, 5, 9), False), + ], +) +@pytest.mark.parametrize("dtype", [np.float32, np.int16, np.complex128]) +def test_getitem_into_numpy_layouts(monkeypatch, shape, chunks, blocks, direct, dtype): + # getitem writes the miniexpr results straight into the NumPy result when the chunks tile it + import importlib + + lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") + monkeypatch.setattr(lazyexpr_mod, "try_miniexpr", True) + assert lazyexpr_mod._chunks_tile_flat(shape, chunks, blocks) is direct + rng = np.random.default_rng(1) + na = (rng.random(shape) * 100).astype(dtype) + nb = (rng.random(shape) * 10).astype(dtype) + a = blosc2.asarray(na, chunks=chunks, blocks=blocks) + b = blosc2.asarray(nb, chunks=chunks, blocks=blocks) + if dtype is np.complex128: + second = (a * b, na * nb) + else: + second = ((a - b) > 40, (na - nb) > 40) + for expr, expected in [(a + b * 2, na + nb * 2), second]: + for res in (expr[:], expr[()]): + assert res.dtype == expected.dtype + assert res.flags.c_contiguous + assert res.flags.writeable + np.testing.assert_array_equal(res, expected) + base = res + while isinstance(getattr(base, "base", None), np.ndarray): + base = base.base + # The direct path's buffer holds a header prefix and at most one chunk of padding + chunk_nbytes = math.prod(chunks) * res.itemsize + nchunks = math.prod(-(-s // c) for s, c in zip(shape, chunks, strict=True)) + in_place = base.dtype == np.uint8 and base.nbytes == nchunks * chunk_nbytes + ( + blosc2.EXTENDED_HEADER_LENGTH + ) + if dtype is np.complex128 and not in_place: + continue # miniexpr is not used for complex numbers on some platforms + assert in_place is direct diff --git a/tests/ndarray/test_reductions.py b/tests/ndarray/test_reductions.py index df4951068..9d0a7b2fa 100644 --- a/tests/ndarray/test_reductions.py +++ b/tests/ndarray/test_reductions.py @@ -983,3 +983,79 @@ 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)) + + +@pytest.mark.parametrize( + ("chunks", "blocks"), + [ + ((10, 20, 50), (5, 20, 50)), # decompressed chunks are C-order; partial along two axes + ((10, 20, 50), (1, 4, 50)), + ((10, 20, 50), (5, 10, 25)), # blocks not contiguous in their chunk + ], +) +@pytest.mark.parametrize("axis", [0, 1, 2, (0, 2), None]) +@pytest.mark.parametrize("op", ["sum", "max", "argmin", "cumulative_sum"]) +def test_reduction_axis_chunk_layouts(chunks, blocks, axis, op): + if op in ("argmin", "cumulative_sum") and not (axis is None or isinstance(axis, int)): + pytest.skip("argmin and cumulative_sum take a single axis") + if op == "cumulative_sum" and axis is None: + pytest.skip("cumulative_sum needs an axis for N-D arrays") + rng = np.random.default_rng(3) + na = rng.integers(-1000, 1000, size=(37, 45, 50)).astype(np.float64) + nb = rng.integers(-1000, 1000, size=(37, 45, 50)).astype(np.float64) + a = blosc2.asarray(na, chunks=chunks, blocks=blocks) + b = blosc2.asarray(nb, chunks=chunks, blocks=blocks) + for expr, nexpr in [(a, na), (a + b, na + nb)]: + res = getattr(expr, op)(axis=axis) + # npcumsum: np.cumulative_sum is NumPy >= 2.1 only + expected = (npcumsum if op == "cumulative_sum" else getattr(np, op))(nexpr, axis=axis) + np.testing.assert_array_equal(res[()] if isinstance(res, blosc2.NDArray) else res, expected) + if op in ("sum", "max"): + res = getattr(expr, op)(axis=axis, keepdims=True) + np.testing.assert_array_equal(res, getattr(np, op)(nexpr, axis=axis, keepdims=True))