diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
new file mode 100644
index 0000000..49b7fe2
--- /dev/null
+++ b/.github/workflows/publish.yml
@@ -0,0 +1,55 @@
+name: Publish to PyPI
+
+# Runs when you publish a GitHub Release (Releases -> Draft a new release -> Publish),
+# OR manually via the "Run workflow" button on the Actions tab (workflow_dispatch).
+# Builds the sdist + wheel and uploads them to PyPI via Trusted Publishing (OIDC) -
+# no API token/secret required (configure the trusted publisher on PyPI once, see README).
+on:
+ release:
+ types: [published]
+ workflow_dispatch:
+
+jobs:
+ build:
+ name: Build distributions
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+
+ - uses: actions/setup-python@v5
+ with:
+ python-version: "3.x"
+
+ - name: Build sdist and wheel
+ run: |
+ python -m pip install --upgrade build
+ python -m build
+
+ - name: Check metadata
+ run: |
+ python -m pip install --upgrade twine
+ python -m twine check dist/*
+
+ - uses: actions/upload-artifact@v4
+ with:
+ name: dist
+ path: dist/
+
+ publish:
+ name: Publish to PyPI
+ needs: build
+ runs-on: ubuntu-latest
+ # The GitHub environment must match the one set on the PyPI trusted publisher.
+ environment:
+ name: pypi
+ url: https://pypi.org/project/dyna-zarr/
+ permissions:
+ id-token: write # OIDC token for Trusted Publishing
+ steps:
+ - uses: actions/download-artifact@v4
+ with:
+ name: dist
+ path: dist/
+
+ - name: Publish
+ uses: pypa/gh-action-pypi-publish@release/v1
diff --git a/.gitignore b/.gitignore
index 6198dad..ba394b5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -61,3 +61,7 @@ env/
# Project-specific
_archive/
.coverage
+
+# Non-shipped / local-only (bench + perf report; tests ARE tracked and CI-run)
+benchmarks/
+reports/
diff --git a/README.md b/README.md
index cd047b5..bd15fa5 100644
--- a/README.md
+++ b/README.md
@@ -1,172 +1,182 @@
-# dyna_zarr
+# dyna-zarr
-A lightweight Python library for lazy operations on Zarr arrays.
+A lightweight, dask-free Python library for lazy, memory-bounded operations on large Zarr (and TIFF) arrays, with an optional GPU path.
## Overview
-**dyna_zarr** provides a thin layer on top of [Zarr](https://zarr-python.readthedocs.io/) for lazy and dynamic array processing. It is designed to simplify working with large, multidimensional datasets by enabling memory-efficient, region-wise I/O and computation.
+dyna-zarr is a thin, pull-based array layer over [Zarr](https://zarr-python.readthedocs.io/). Instead of building a task graph, every operation is a lazy *transform* whose `read(key)` maps an output slice back to a bounded input read, ending at a direct zarr/TensorStore read. Slicing a result pulls only that region through the whole operation chain, so no intermediates are materialized.
+
+The practical consequence is memory-boundedness. When you stream a result to disk with `io.write`, the array is processed region by region, so peak RAM is a function of the region and worker budget rather than the array size. This makes it possible to read, transform, and write arrays far larger than memory.
+
+## Memory-boundedness
+
+There are two ways to run a lazy result, with different memory behavior:
+
+- `io.write(result, path)` streams the result to disk region by region. Peak RAM is roughly `region_size_mb * max_workers`, independent of the array size. This is the memory-bounded path.
+- `result.compute()` returns a single in-memory NumPy array. It materializes the whole result by design (mirroring `dask.array.compute`), so it is not memory-bounded. Use it only for results that fit in RAM.
+
+Every operation is memory-bounded on the `io.write` path except `median`, `argmin`, and `argmax`, which are flagged in the operations catalog below.
## Features
-- **Lazy evaluation** – Operations are deferred until explicitly computed
-- **Multi-format I/O** – Read from TIFF, Zarr v2, and Zarr v3; write to Zarr v2 or Zarr v3
-- **Efficient region-wise processing** – Data is processed in regions, where each region may span multiple chunks
-- **Minimal dependencies** – Requires only `zarr`, `numpy`, `tensorstore`, and `tifffile`
+- **Pull-based and lazy.** Operations defer until `.compute()` (materialize) or `io.write` (stream to disk).
+- **Memory-bounded streaming.** Region-wise `io.write` with per-worker memory and worker-count knobs. Even reshape, flatten, and rechunk of incompatibly-chunked data stay bounded, by staging through disk.
+- **NumPy-like.** Operator overloads, array methods (`.astype`, `.clip`, `.round`), and the NumPy ufunc protocol (`np.sqrt(a)`, `np.add(a, 2)`) all work on a `DynamicArray`.
+- **Rich op set.** About 90 operations: pointwise ufuncs, streaming reductions, neighborhood (halo) filters, structural reshaping, differences, and array creation.
+- **Multi-format I/O.** Read TIFF, Zarr v2, and Zarr v3 (local, S3/GCS, HTTP); write Zarr v2/v3 with optional sharding.
+- **Optional GPU.** Run an op chain on CUDA via CuPy, with a single host-to-device transfer per region.
## Installation
```bash
-pip install git+https://github.com/bugraoezdemir/dyna_zarr.git
+pip install dyna-zarr
```
-For development:
+Optional GPU support (pick the extra matching your CUDA toolkit from `nvidia-smi`):
```bash
-git clone https://github.com/bugraoezdemir/dyna_zarr.git
-cd dyna_zarr
-pip install -e ".[dev]"
+pip install "dyna-zarr[gpu-cu12]" # CUDA 12.x ([gpu] is an alias for this)
+pip install "dyna-zarr[gpu-cu11]" # CUDA 11.x
+pip install "dyna-zarr[gpu-cu13]" # CUDA 13.x (e.g. Blackwell)
```
-## Quick Start
+## Quick start
-### Reading Arrays
+### Read
```python
from dyna_zarr import io
-# Read a TIFF file as a DynamicArray
-arr = io.read("image.tiff")
-print(arr.shape) # (100, 256, 256)
-print(arr.dtype) # dtype('uint16')
-
-# Read Zarr v2
-arr = io.read("array_v2.zarr")
-
-# Read Zarr v3
-arr = io.read("array_v3.zarr")
+arr = io.read("image.tiff") # TIFF via tifffile's zarr bridge
+arr = io.read("array_v2.zarr") # Zarr v2
+arr = io.read("array_v3.zarr") # Zarr v3 (also s3://, gs://, http://)
-# Compute the full array into memory
-data = arr.compute()
+print(arr.shape, arr.dtype, arr.chunks)
-# Compute a region of interest
-region = arr[10:20, 50:150, 100:200].compute()
+data = arr.compute() # materialize the whole array
+region = arr[10:20, 50:150, 100:200].compute() # pull just this region
```
-### Writing Arrays
+### Write (memory-bounded streaming)
```python
from dyna_zarr import io, Codecs
-# Write to Zarr v3
-io.write(arr, "output_v3.zarr", zarr_format=3)
+io.write(arr, "out_v3.zarr", zarr_format=3)
+io.write(arr, "out.zarr", chunks=(64, 64, 64), zarr_format=3)
+io.write(arr, "out.zarr", dtype="float32", zarr_format=3) # cast on write
+io.write(arr, "out.zarr", compressor=Codecs(compressor="zstd", clevel=5), zarr_format=3)
-# Write to Zarr v2
-io.write(arr, "output_v2.zarr", zarr_format=2)
-
-# Specify custom chunks
-io.write(arr, "output.zarr", chunks=(64, 64, 64), zarr_format=3)
-
-# Enable compression
-codecs = Codecs(compressor="zstd", clevel=5)
-io.write(arr, "output.zarr", compressor=codecs, zarr_format=3)
-
-# Convert dtype during write
-io.write(arr, "output.zarr", dtype="float32", zarr_format=3)
+# memory and parallelism controls (peak RAM is roughly region_size_mb * max_workers)
+io.write(arr, "out.zarr", region_size_mb=64, max_workers=4)
```
-### Lazy Operations
+### Lazy operation chains
```python
-from dyna_zarr import io, operations
+from dyna_zarr import io, operations as ops
-# Read array
arr = io.read("input.zarr")
-# Apply lazy transformations (no computation yet)
-result = operations.abs(arr)
-result = operations.clip(result, 0, 1)
-result = operations.sqrt(result)
+result = ops.sqrt(ops.clip(ops.abs(arr), 0, 1)) # nothing computed yet
+io.write(result, "output.zarr", zarr_format=3) # streamed, region by region
+# ...or result.compute() to materialize
+```
-# Write result using region-wise processing
-io.write(result, "output.zarr", zarr_format=3)
+### NumPy-like interface
+
+A `DynamicArray` behaves like a NumPy or dask array. Operators, methods, and ufuncs are all lazy:
+
+```python
+import numpy as np
-# Or fully materialize the result
-final_data = result.compute()
+masked = (arr > 3) & (arr < 100) # elementwise operators build a lazy mask
+scaled = (arr.astype("float32") / 255).clip(0, 1)
+out = np.sqrt(np.abs(arr)) # NumPy ufunc protocol dispatches to lazy ops
```
-### Multi-source Operations
+### Neighborhood filters
+
+Neighborhood (halo) filters wrap `scipy.ndimage`. Each read pulls its own halo, so results are chunk-invariant and exact, and stay memory-bounded when streamed.
```python
-from dyna_zarr import io, operations
+import numpy as np
+from dyna_zarr import io, operations as ops
+
+img = io.read("volume.zarr") # e.g. (z, y, x)
-# Read multiple sources
-arr1 = io.read("input1.zarr")
-arr2 = io.read("input2.zarr")
+# LoG filtering
+log = ops.gaussian_laplace(img, sigma=2)
+io.write(log, "log.zarr", zarr_format=3) # halo handled per region
-# Chain operations
-result = operations.concatenate([arr1, arr2], axis=0)
-result = operations.clip(result, -1, 1)
+# median denoise
+denoised = ops.median_filter(img, size=3)
+io.write(denoised, "denoised.zarr")
-# Write result
-io.write(result, "concatenated_output.zarr", zarr_format=3)
+# a custom per-plane kernel
+kernel = np.ones((1, 3, 3), dtype="float32") / 9 # 3x3 mean within each z-plane
+blurred = ops.convolve(img, kernel)
+io.write(blurred, "blurred.zarr")
```
-## Core Components
+## Operations catalog
-- **`io.read()`** – Read TIFF, Zarr v2, or Zarr v3 sources and return a
- `DynamicArray`
-- **`io.write()`** – Write a `DynamicArray` to Zarr v2 or v3 with optional
- region-wise execution
-- **`operations`** – Lazy transformation functions such as `abs`,
- `clip`, `sqrt`, `concatenate`, `reshape`, and `slice`
-- **`DynamicArray`** – Core lazy array abstraction supporting slicing,
- shape/dtype inspection, and `.compute()`
-- **`Codecs`** – Compression configuration for Zarr v2 and v3
+Every operation is lazy, and memory-bounded on the `io.write` path except `median`, `argmin`, and `argmax` (see Memory-boundedness). All are available flat on `dyna_zarr.operations`, and also grouped by category submodule.
-## Further Examples
+- **Pointwise / ufuncs.** `abs`, `negative`, `sign`, `sqrt`, `square`, `exp`, `log`, `log2`, `log10`, `floor`, `ceil`, `reciprocal`, `round`, `clip`, `astype`; binary `add`, `subtract`, `multiply`, `divide`, `floor_divide`, `mod`, `power`, `maximum`, `minimum`; comparisons `greater(_equal)`, `less(_equal)`, `equal`, `not_equal`; logical `and`, `or`, `xor`, `not`; `where`, `isin`, `digitize`.
+- **Reductions.** Streaming and memory-bounded: `min`, `max`, `sum`, `prod`, `mean`, `any`, `all`, `var`, `std`, `histogram` (with `axis=` and `keepdims=`). Not fully bounded (hold the full reduced axis): `median`, `argmin`, `argmax`.
+- **Neighborhood (halo/overlap).** `gaussian_filter`, `uniform_filter`, `median_filter`, `minimum_filter`, `maximum_filter`, `grey_erosion`, `grey_dilation`, `convolve`, `correlate`, `laplace`, `gaussian_laplace`, `gaussian_gradient_magnitude`.
+- **Structural.** `concatenate`, `stack`, `transpose`, `swap_axes`, `reshape`, `flatten`, `squeeze`, `expand_dims`, `pad`, `tile`, `roll`, `flip`, `rot90`, `slice_array`.
+- **Differences.** `diff`, `gradient`.
+- **Scan (prefix, along one axis).** `cumsum`, `cumprod`, `cummax`, `cummin`. Streamed with a bounded carry on the `io.write` path, so memory-bounded despite the sequential dependency.
+- **Creation.** `zeros`, `ones`, `full`, `empty`, `random` (and the `*_like` variants). `random` is position-deterministic, so the result is independent of chunking.
+- **Primitives.** `map_blocks` (pointwise), `map_overlap` (neighborhood with a halo), `reduce` (streaming). Use these to build your own ops.
-### Manual Configuration of Region Size
+## Memory-bounded reshape, flatten, and rechunk
-```python
-from dyna_zarr import io
+C-order reshape and flatten conflict with n-dimensional chunk layout, so a naive implementation blows up. dyna-zarr stages these through disk (a Rechunker-style two-phase, read-once/write-once copy), so peak RAM stays a function of the per-worker budget rather than the array size. When you `io.write` an outermost `reshape` or `flatten`, this path is used automatically:
-arr = io.read("large_array.zarr")
+```python
+from dyna_zarr import io, operations as ops
-# Process approximately 64 MB regions at a time
-io.write(arr, "output.zarr", region_size_mb=64)
+arr = io.read("big_4d.zarr") # e.g. 5 GB, awkward chunks
+io.write(ops.flatten(arr), "flat.zarr", region_size_mb=128, max_workers=2)
+io.write(ops.reshape(arr, (a, b)), "reshaped.zarr") # (a, b) is any target shape of the same size
```
-### Zarr Sharding (Zarr v3 only)
+## GPU (optional)
+
+With a CuPy install, run a chain on the GPU. Setting `device='cuda'` on a terminal call (`compute` or `io.write`) makes device-inheriting ops run on the GPU. A single host-to-device transfer happens at the first CUDA op and the data stays resident up the chain. Results are returned or written from the host.
```python
-io.write(
- arr,
- "sharded_output.zarr",
- chunks=(64, 64, 64), # Inner chunk size
- shard_coefficients=(4, 4, 4), # Shard size = 4x chunks in each dimension
- zarr_format=3,
-)
+result = ops.gaussian_filter(arr, sigma=3)
+out = result.compute(device="cuda") # whole chain on the GPU
+io.write(result, "out.zarr", device="cuda") # per-region GPU compute, streamed write
```
-## Requirements
+## Relationship to dask
-- Python >= 3.11
-- zarr >= 3.0.0
-- numpy >= 1.20.0
-- tensorstore
-- tifffile
+dyna-zarr is not a general replacement for `dask.array`. It targets one job: memory-bounded read, transform, and write of large Zarr/TIFF arrays.
-## Testing
+The core idea is to drop the task graph. Because every operation is a pull-based chain, where each output slice maps back to a bounded input read, there is no graph to build and no scheduler to run it. That keeps the engine small, keeps peak RAM bounded by `region_size_mb * max_workers` on the `io.write` path, and avoids scheduling overhead, which makes the read-transform-write pipeline efficient.
-Run the test suite:
+The tradeoff is that only operations that fit this slice-pushdown model belong in the chain: pointwise math, neighborhood/halo filters, streaming reductions, and structural reshaping. These are operations that are commonly used in image processing, which is what dyna-zarr is mainly built for. Operations that would need a global, data-dependent graph do not fit directly, and a few that do (such as non-associative reductions) trade extra reads or memory to stay correct.
-```bash
-pytest tests/ -v
-```
+Two more differences worth knowing:
-Run tests with coverage:
+- **Single machine, for now.** Parallelism today is threaded I/O within one process, plus the optional GPU path. There is no cluster or distributed execution yet; better and process-based parallelism is a possible future direction.
+- **Narrower surface.** About 90 operations today, extended where the slice-pushdown model permits. Binary ops also need equal-shaped operands (no general broadcasting between differently shaped lazy arrays yet).
-```bash
-pytest tests/ --cov=src/dyna_zarr --cov-report=html
-```
+## Core components
+
+- `io.read(source)` reads TIFF, Zarr v2, or Zarr v3 (local or remote) into a `DynamicArray`.
+- `io.write(array, path, ...)` streams a `DynamicArray` to Zarr v2/v3 (chunks, sharding, compression, dtype cast, `region_size_mb`, `max_workers`, `device`).
+- `operations` is the lazy op set above.
+- `DynamicArray` is the pull-based lazy array (slicing, `.compute()`, operators, `.astype`/`.clip`/`.round`, ufunc protocol).
+- `Codecs` is the compression configuration for Zarr v2 and v3.
+
+## Requirements
+- Python 3.11 or newer
+- zarr 3.0.0+, numpy 1.20+, scipy 1.6+, tensorstore, tifffile
+- Optional: CuPy (via the `gpu-cuXX` extras) for the GPU path
diff --git a/pyproject.toml b/pyproject.toml
index 4336fd0..51a8dc8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,8 +9,8 @@ packages = ["dyna_zarr"]
"" = "src"
[project]
-name = "dyna_zarr"
-version = "0.0.1"
+name = "dyna-zarr"
+version = "0.0.2"
description = "A lightweight library for lazy operations on Zarr arrays without task graph overhead"
readme = "README.md"
requires-python = ">=3.11"
@@ -32,6 +32,7 @@ classifiers = [
dependencies = [
"zarr>=3.0.0",
"numpy>=1.20.0",
+ "scipy>=1.6.0",
"tensorstore",
"tifffile"
]
@@ -49,6 +50,18 @@ docs = [
"sphinx>=5.0",
"sphinx-rtd-theme>=1.0",
]
+# GPU execution (optional). CuPy ships CUDA-version-specific wheels and pip CANNOT
+# auto-detect your CUDA, so pick the extra that matches `nvidia-smi` (CUDA Version):
+# pip install dyna-zarr[gpu-cu11] # CUDA 11.x
+# pip install dyna-zarr[gpu-cu12] # CUDA 12.x
+# pip install dyna-zarr[gpu-cu13] # CUDA 13.x (e.g. Blackwell)
+# `[gpu]` is a convenience alias for CUDA 12.x. The [ctk] extra ships the NVRTC headers
+# CuPy needs to JIT kernels (without it every kernel raises "Failed to find CUDA headers").
+# Auto-detect alternative (may lag new CUDA releases): pip install cupy-wheel
+gpu = ["cupy-cuda12x[ctk]"]
+gpu-cu11 = ["cupy-cuda11x[ctk]"]
+gpu-cu12 = ["cupy-cuda12x[ctk]"]
+gpu-cu13 = ["cupy-cuda13x[ctk]"]
[project.urls]
Homepage = "https://github.com/bugraoezdemir/dyna_zarr"
diff --git a/scripts/clean_extracted.py b/scripts/clean_extracted.py
deleted file mode 100644
index 24a348b..0000000
--- a/scripts/clean_extracted.py
+++ /dev/null
@@ -1,20 +0,0 @@
-from pathlib import Path
-p = Path('htmlcov_extracted_dynamic_array_fixed.py')
-text = p.read_text()
-lines = text.splitlines()
-# find first line that starts with 'import zarr'
-start = 0
-for i,l in enumerate(lines):
- if l.strip().startswith('import zarr'):
- start = i
- break
-# find last occurrence of the write return
-end = 0
-for i in range(len(lines)-1, -1, -1):
- if 'return result._with_transform(transform)' in lines[i] or 'return array._with_transform(transform)' in lines[i]:
- end = i
- break
-clean = '\n'.join(lines[start:end+1])
-out = Path('htmlcov_extracted_dynamic_array_clean.py')
-out.write_text(clean)
-print('wrote', out)
diff --git a/scripts/extract_dynamic_array.py b/scripts/extract_dynamic_array.py
deleted file mode 100644
index 7f9620f..0000000
--- a/scripts/extract_dynamic_array.py
+++ /dev/null
@@ -1,22 +0,0 @@
-import re
-import html
-from pathlib import Path
-html_path = Path('htmlcov/z_5bb88312f9bf08c8_dynamic_array_py.html')
-out_path = Path('htmlcov_extracted_dynamic_array.py')
-text = html_path.read_text()
-# extract all ...
-parts = re.findall(r']*>(.*?)', text, flags=re.S)
-lines = []
-for p in parts:
- # remove tags
- s = re.sub(r'<[^>]+>', '', p)
- s = html.unescape(s)
- # remove leading line numbers if present
- s = re.sub(r'^\s*\d+\s*', '', s)
- # strip trailing non-breaking spaces
- s = s.replace('\xa0', ' ')
- lines.append(s.rstrip())
-# join and write
-out = '\n'.join(lines)
-out_path.write_text(out)
-print('wrote', out_path)
diff --git a/scripts/extract_dynamic_array_fix.py b/scripts/extract_dynamic_array_fix.py
deleted file mode 100644
index 9027a4f..0000000
--- a/scripts/extract_dynamic_array_fix.py
+++ /dev/null
@@ -1,18 +0,0 @@
-import re
-import html
-from pathlib import Path
-html_path = Path('htmlcov/z_5bb88312f9bf08c8_dynamic_array_py.html')
-out_path = Path('htmlcov_extracted_dynamic_array_fixed.py')
-text = html_path.read_text()
-parts = re.findall(r']*>(.*?)', text, flags=re.S)
-lines = []
-for p in parts:
- s = re.sub(r'<[^>]+>', '', p)
- s = html.unescape(s)
- # remove only the leading line number, keep following spaces
- s = re.sub(r'^\s*\d+', '', s)
- s = s.replace('\xa0', ' ')
- lines.append(s.rstrip())
-out = '\n'.join(lines)
-out_path.write_text(out)
-print('wrote', out_path)
diff --git a/scripts/extract_grid_ok.py b/scripts/extract_grid_ok.py
deleted file mode 100644
index 1aba309..0000000
--- a/scripts/extract_grid_ok.py
+++ /dev/null
@@ -1,57 +0,0 @@
-#!/usr/bin/env python3
-import json
-import csv
-from pathlib import Path
-
-ROOT = Path(__file__).resolve().parents[1]
-JSON_PATH = ROOT / 'grid_search_results.json'
-OUT_CSV = ROOT / 'grid_search_ok_sorted.csv'
-
-with JSON_PATH.open() as f:
- data = json.load(f)
-
-ok_entries = []
-config_keys = set()
-for i, entry in enumerate(data):
- res = entry.get('result', {})
- if res.get('status') == 'ok':
- cfg = entry.get('config', {})
- config_keys.update(cfg.keys())
- row = {
- 'index': i,
- 'elapsed': res.get('elapsed'),
- 'trial_time': entry.get('trial_time')
- }
- # flatten config
- for k, v in cfg.items():
- row[f'cfg_{k}'] = v
- # include any other result fields except status/elapsed
- for k, v in res.items():
- if k in ('status','elapsed'):
- continue
- row[f'result_{k}'] = v
- ok_entries.append(row)
-
-# sort by elapsed (None goes to end)
-ok_entries.sort(key=lambda r: (r['elapsed'] is None, r['elapsed']))
-
-# build header
-cfg_keys_sorted = sorted(list(config_keys))
-headers = ['index','elapsed','trial_time'] + [f'cfg_{k}' for k in cfg_keys_sorted]
-# collect extra result keys
-extra_result_keys = set()
-for r in ok_entries:
- for k in r.keys():
- if k.startswith('result_'):
- extra_result_keys.add(k)
-headers += sorted(extra_result_keys)
-
-with OUT_CSV.open('w', newline='') as csvfile:
- writer = csv.DictWriter(csvfile, fieldnames=headers)
- writer.writeheader()
- for r in ok_entries:
- # ensure all header keys exist
- out = {h: r.get(h, '') for h in headers}
- writer.writerow(out)
-
-print(f'Wrote {len(ok_entries)} successful trials to {OUT_CSV}')
diff --git a/scripts/extract_grid_search_results.py b/scripts/extract_grid_search_results.py
deleted file mode 100644
index cc65488..0000000
--- a/scripts/extract_grid_search_results.py
+++ /dev/null
@@ -1,57 +0,0 @@
-#!/usr/bin/env python3
-"""Extract successful trials from grid_search_results.json and write CSV.
-
-Usage: python3 scripts/extract_grid_search_results.py
-Creates: grid_search_ok_sorted.csv in repository root.
-"""
-import json
-import csv
-from pathlib import Path
-
-ROOT = Path(__file__).resolve().parent.parent
-SRC = ROOT / "grid_search_results.json"
-OUT = ROOT / "grid_search_ok_sorted.csv"
-
-
-def main():
- data = json.loads(SRC.read_text())
-
- # collect all config keys to build CSV header
- config_keys = set()
- rows = []
- for entry in data:
- result = entry.get("result", {})
- if result.get("status") != "ok":
- continue
- config = entry.get("config", {})
- config_keys.update(config.keys())
- elapsed = result.get("elapsed")
- trial_time = entry.get("trial_time")
- rows.append({"elapsed": elapsed, "trial_time": trial_time, **config})
-
- if not rows:
- print("No successful trials found.")
- return
-
- # order config columns consistently
- config_keys = sorted(config_keys)
-
- # sort rows by elapsed (ascending)
- rows.sort(key=lambda r: float(r.get("elapsed", float("inf"))))
-
- # write CSV header
- header = ["elapsed", "trial_time"] + config_keys
-
- with OUT.open("w", newline="") as f:
- w = csv.DictWriter(f, fieldnames=header)
- w.writeheader()
- for r in rows:
- # ensure all keys present
- row = {k: r.get(k, "") for k in header}
- w.writerow(row)
-
- print(f"Wrote {len(rows)} successful trials to {OUT}")
-
-
-if __name__ == "__main__":
- main()
diff --git a/src/dyna_zarr/__init__.py b/src/dyna_zarr/__init__.py
index 4fc6ee0..20abac0 100644
--- a/src/dyna_zarr/__init__.py
+++ b/src/dyna_zarr/__init__.py
@@ -8,21 +8,21 @@
Includes efficient TIFF reading via tifffile's zarr bridge with concurrent access support.
"""
-__version__ = "0.0.1"
+__version__ = "0.0.2"
__author__ = "EuBI-Biohub"
# Import core classes
from .dynamic_array import DynamicArray, slice_array
from .codecs import Codecs
-# Import operations and io namespaces
-from .operations import operations
+# Import operations (now a package/module of flat ops) and io namespaces
+from . import operations
from .io import io
# Optional: expose tifffile utilities if available
try:
from .io import read_file
- from .tiff_reader import read_tiff_lazy, TiffZarrReader
+ from .tiff_reader import read_tiff_lazy, open_tiff_zarr
__all__ = [
"DynamicArray",
"operations",
@@ -31,7 +31,7 @@
"Codecs",
"read_file",
"read_tiff_lazy",
- "TiffZarrReader",
+ "open_tiff_zarr",
]
except ImportError:
__all__ = [
diff --git a/src/dyna_zarr/codecs.py b/src/dyna_zarr/codecs.py
index 76262bd..1d40219 100644
--- a/src/dyna_zarr/codecs.py
+++ b/src/dyna_zarr/codecs.py
@@ -108,54 +108,54 @@ def to_numcodecs(self):
def to_v3_config(self):
"""
- Generate codec pipeline for Zarr v3.
+ Generate codec pipeline for Zarr v3 using proper zarr.codecs API.
Returns:
- List of codec configurations
+ List of codec configurations (from .to_dict() calls)
"""
- codecs = [
- {'name': 'bytes', 'configuration': {'endian': 'little'}}
+ from zarr import codecs
+
+ codecs_list = [
+ codecs.BytesCodec(endian=codecs.Endian.little).to_dict()
]
if self.compressor == 'blosc':
- # Convert shuffle integer to string for TensorStore
- shuffle_map = {0: 'noshuffle', 1: 'shuffle', 2: 'bitshuffle'}
- shuffle_str = shuffle_map.get(self.shuffle, 'shuffle')
+ # Use BloscCodec with proper BloscShuffle enum
+ shuffle_map = {0: codecs.BloscShuffle.noshuffle,
+ 1: codecs.BloscShuffle.shuffle,
+ 2: codecs.BloscShuffle.bitshuffle}
+ shuffle_enum = shuffle_map.get(self.shuffle, codecs.BloscShuffle.shuffle)
+
+ blosc_codec = codecs.BloscCodec(
+ cname=self.cname,
+ clevel=self.clevel,
+ shuffle=shuffle_enum
+ )
+ codecs_list.append(blosc_codec.to_dict())
- codecs.append({
- 'name': 'blosc',
- 'configuration': {
- 'cname': self.cname,
- 'clevel': self.clevel,
- 'shuffle': shuffle_str
- }
- })
elif self.compressor == 'zstd':
- codecs.append({
- 'name': 'zstd',
- 'configuration': {'level': self.clevel}
- })
+ zstd_codec = codecs.ZstdCodec(level=self.clevel)
+ codecs_list.append(zstd_codec.to_dict())
+
elif self.compressor == 'gzip':
- codecs.append({
- 'name': 'gzip',
- 'configuration': {'level': self.clevel}
- })
+ gzip_codec = codecs.GzipCodec(level=self.clevel)
+ codecs_list.append(gzip_codec.to_dict())
+
elif self.compressor == 'lz4':
- codecs.append({
- 'name': 'lz4',
- 'configuration': {}
- })
+ # LZ4 is not directly available in zarr.codecs v3, would need blosc
+ raise ValueError("LZ4 is not available in Zarr v3. Use Blosc with cname='lz4' instead.")
+
elif self.compressor == 'bz2':
- codecs.append({
- 'name': 'bz2',
- 'configuration': {'level': self.clevel}
- })
+ # BZ2 is not directly available in zarr.codecs v3
+ raise ValueError("BZ2 is not available in Zarr v3. Use Zstd or Gzip instead.")
+
elif self.compressor is None or self.compressor == 'none':
- pass # No compression codec
+ pass # No compression codec, only bytes codec
+
else:
raise ValueError(f"Unsupported compressor: {self.compressor}")
- return codecs
+ return codecs_list
@classmethod
def from_numcodecs(cls, compressor):
diff --git a/src/dyna_zarr/dynamic_array.py b/src/dyna_zarr/dynamic_array.py
index 5e73ac4..62aaa7a 100755
--- a/src/dyna_zarr/dynamic_array.py
+++ b/src/dyna_zarr/dynamic_array.py
@@ -261,21 +261,112 @@ def __getitem__(self, key):
transform = SliceTransform(self, key)
return self._with_transform(transform)
- def compute(self):
+ def compute(self, device=None):
"""
Execute all lazy transforms and return the result as a numpy array.
+
+ ``device`` sets the EXECUTION device for inherit-ops (device=None) in the chain:
+ None/'cpu' runs on the CPU (default), 'cuda' runs the pipeline on the GPU. The
+ result is always returned as a host numpy array.
"""
- if self._transform is None:
- # No transform - read entire array
- if self._is_tensorstore:
- return self._ts_array[:].read().result()
+ from .operations._backend import device_context, to_device
+ with device_context(device):
+ if self._transform is None:
+ # No transform - read entire array
+ if self._is_tensorstore:
+ result = self._ts_array[:].read().result()
+ else:
+ result = self._zarr_array[:]
else:
- return self._zarr_array[:]
- else:
- # Apply transformation to read all data
- full_slice = tuple(slice(None) for _ in range(len(self.shape)))
- return self._transform.read(full_slice)
-
+ # Apply transformation to read all data
+ full_slice = tuple(slice(None) for _ in range(len(self.shape)))
+ result = self._transform.read(full_slice)
+ return to_device(result, "cpu")
+
+ # --- numpy/dask-like method surface (lazy; route to operations) ---
+ def astype(self, dtype):
+ """Lazily cast to ``dtype`` (like ``numpy``/``dask`` ``a.astype``)."""
+ from . import operations as o
+ return o.astype(self, dtype)
+
+ def clip(self, a_min=None, a_max=None, out=None):
+ """Lazily clip to ``[a_min, a_max]`` (either bound may be None). ``out`` accepted for
+ NumPy-method compatibility (``np.clip`` calls ``a.clip(min, max, out=...)``); a non-None
+ ``out`` isn't supported on a lazy array."""
+ if out is not None:
+ raise TypeError("clip(out=...) is not supported on a lazy DynamicArray")
+ from . import operations as o
+ return o.clip(self, a_min, a_max)
+
+ def round(self, decimals=0, out=None):
+ """Lazily round (like ``numpy``/``dask`` ``a.round``). ``out`` accepted for NumPy-method
+ compatibility; a non-None ``out`` isn't supported on a lazy array."""
+ if out is not None:
+ raise TypeError("round(out=...) is not supported on a lazy DynamicArray")
+ from . import operations as o
+ return o.round(self, decimals)
+
+ def rechunk(self, chunks=None, **kwargs):
+ """No-op for the pull model (accepts dask's signature for backend compatibility).
+
+ In dask, ``rechunk`` changes the chunk grid so cross-chunk ops behave. A DynamicArray
+ is chunk-invariant: every lazy read pulls exactly the region asked for, independent of
+ any chunk grid, and streaming reductions/scans already see whole axes. So a lazy
+ rechunk changes nothing about correctness and returns the array unchanged. Storage
+ chunking is a separate concern, set via ``io.write(chunks=...)`` or the rechunk engine.
+ """
+ return self
+
+ def persist(self, **kwargs):
+ """No-op (accepts dask's signature). dask ``persist`` materializes and caches an
+ intermediate; the pull model has no graph to cache, so this returns the array
+ unchanged. Use ``io.write`` to stage an intermediate to disk when needed."""
+ return self
+
+ def map_blocks(self, func, *args, dtype=None, **kwargs):
+ """dask-compatible ``map_blocks``: apply ``func`` blockwise (shape-preserving). Extra
+ array/scalar ``args`` become additional equally-shaped operands; dask-only kwargs
+ (``meta``/``chunks``/``name``/``block_info``/...) are ignored, and any remaining kwargs
+ are bound to ``func``. ``drop_axis``/``new_axis`` (shape-changing) are not supported."""
+ from . import operations as o
+ if kwargs.get("drop_axis") is not None or kwargs.get("new_axis") is not None:
+ raise NotImplementedError(
+ "map_blocks drop_axis/new_axis is not supported (shape-preserving only)")
+ for k in ("meta", "chunks", "name", "token", "drop_axis", "new_axis",
+ "block_info", "block_id", "enforce_ndim"):
+ kwargs.pop(k, None)
+ f = (lambda *bs, _f=func, _kw=kwargs: _f(*bs, **_kw)) if kwargs else func
+ return o.map_blocks(f, self, *args, dtype=dtype)
+
+ def map_overlap(self, func, depth=0, boundary="reflect", trim=True, dtype=None, **kwargs):
+ """dask-compatible ``map_overlap``: apply a shape-preserving neighbourhood ``func``
+ with a ``depth`` halo. dask-only kwargs are ignored and any remaining kwargs are bound
+ to ``func``. ``trim=False`` (shrinking output) is not supported."""
+ from . import operations as o
+ if not trim:
+ raise NotImplementedError("map_overlap trim=False is not supported")
+ for k in ("meta", "chunks", "name"):
+ kwargs.pop(k, None)
+ f = (lambda b, _f=func, _kw=kwargs: _f(b, **_kw)) if kwargs else func
+ return o.map_overlap(self, f, depth, boundary=boundary, dtype=dtype)
+
+ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+ """NumPy ufunc protocol -> lazy ops, so ``np.sqrt(a)`` / ``np.add(a, 2)`` work on a
+ DynamicArray exactly as on a dask array. Only the plain ``__call__`` form (no ``out=``,
+ no reductions) is handled; anything else -- or a ufunc dyna_zarr doesn't implement --
+ returns ``NotImplemented`` so NumPy can fall back / raise clearly."""
+ from . import operations as o
+ if method != "__call__" or kwargs.get("out") is not None:
+ return NotImplemented
+ name = _UFUNC_ALIASES.get(ufunc.__name__, ufunc.__name__)
+ fn = getattr(o, name, None)
+ if fn is None:
+ return NotImplemented
+ try:
+ return fn(*inputs)
+ except (TypeError, ValueError):
+ return NotImplemented
+
def _read_direct(self, key):
"""
Internal method to read data directly without creating transforms.
@@ -304,46 +395,127 @@ def _with_transform(self, transform):
result._dtype = transform.dtype
# Keep zarr metadata from original
return result
+
+ @classmethod
+ def _from_transform(cls, transform):
+ """Build a SOURCELESS DynamicArray whose data is synthesized by ``transform`` (a
+ generative source, e.g. the creation ops). There is no underlying zarr/tensorstore
+ array; every read goes through ``transform.read(key)``."""
+ self = cls.__new__(cls)
+ self._source = None
+ self._zarr_array = None
+ self._ts_array = None
+ self._is_tensorstore = False
+ self._shape = tuple(transform.shape)
+ self._chunks = transform.chunks
+ self._dtype = np.dtype(transform.dtype)
+ self._transform = transform
+ self._zarr_format = None
+ self._compressor = None
+ self._compressors = None
+ self._shards = None
+ self._codecs = None
+ return self
- def min(self, axis: Optional[int] = None):
- """Compute minimum along specified axis.
-
- Parameters
- ----------
- axis : int, optional
- Axis along which to compute minimum. If None, reduces to scalar.
-
- Returns
- -------
- scalar or numpy.ndarray
- Minimum value(s). Computed immediately.
- """
+ # Eager reduction shortcuts. Each streams via operations. (memory-bound)
+ # and computes immediately, returning a NumPy scalar/array. Use operations.
+ # directly (e.g. operations.max(a, 0)) for a lazy, chainable/writable reduction.
+ def min(self, axis=None, keepdims=False):
+ """Minimum along ``axis`` (None = all). Computed immediately."""
from . import operations
- result = operations.min(self, axis=axis)
- return result.compute()
-
- def max(self, axis: Optional[int] = None):
- """Compute maximum along specified axis.
-
- Parameters
- ----------
- axis : int, optional
- Axis along which to compute maximum. If None, reduces to scalar.
-
- Returns
- -------
- scalar or numpy.ndarray
- Maximum value(s). Computed immediately.
- """
+ return operations.min(self, axis=axis, keepdims=keepdims).compute()
+
+ def max(self, axis=None, keepdims=False):
+ """Maximum along ``axis`` (None = all). Computed immediately."""
+ from . import operations
+ return operations.max(self, axis=axis, keepdims=keepdims).compute()
+
+ def sum(self, axis=None, keepdims=False):
+ """Sum along ``axis`` (None = all). Computed immediately."""
+ from . import operations
+ return operations.sum(self, axis=axis, keepdims=keepdims).compute()
+
+ def mean(self, axis=None, keepdims=False):
+ """Mean along ``axis`` (None = all). Computed immediately."""
+ from . import operations
+ return operations.mean(self, axis=axis, keepdims=keepdims).compute()
+
+ def prod(self, axis=None, keepdims=False):
+ """Product along ``axis`` (None = all). Computed immediately."""
+ from . import operations
+ return operations.prod(self, axis=axis, keepdims=keepdims).compute()
+
+ def median(self, axis=None, keepdims=False):
+ """Median along ``axis`` (None = all). Computed immediately."""
from . import operations
- result = operations.max(self, axis=axis)
- return result.compute()
+ return operations.median(self, axis=axis, keepdims=keepdims).compute()
+
+ def histogram(self, bins=256, range=None):
+ """Streaming histogram over the whole array. Returns ``(counts, bin_edges)`` like
+ numpy. Slice first for a per-channel/plane histogram: ``da[channel].histogram()``."""
+ from . import operations
+ return operations.histogram(self, bins=bins, range=range)
# Import Transform subclasses
from .operations import SliceTransform
+# --------------------------------------------------------------------------- #
+# Operator overloads -> lazy elementwise ops via operations.map_blocks.
+# Mirrors numpy / ome_zarr_pyramid.Pyramid: arithmetic and comparisons are
+# elementwise; &,|,^,~ are LOGICAL (for boolean masks), matching `~mask`. `a == b`
+# returns a lazy mask, but assigning __eq__ *after* the class body leaves the
+# inherited identity __hash__ intact, so a DynamicArray is still hashable.
+# --------------------------------------------------------------------------- #
+
+# numpy ufunc names that differ from the operations spelling (see __array_ufunc__).
+_UFUNC_ALIASES = {
+ "absolute": "abs",
+ "true_divide": "divide",
+ "bitwise_and": "logical_and",
+ "bitwise_or": "logical_or",
+ "bitwise_xor": "logical_xor",
+ "invert": "logical_not",
+}
+
+
+def _binop(opname, reflected=False):
+ def method(self, other):
+ from . import operations as o
+ fn = getattr(o, opname)
+ return fn(other, self) if reflected else fn(self, other)
+ method.__name__ = ("__r" if reflected else "__") + opname + "__"
+ return method
+
+
+def _unop(opname):
+ def method(self):
+ from . import operations as o
+ return getattr(o, opname)(self)
+ method.__name__ = "__" + opname + "__"
+ return method
+
+
+for _dunder, _op in {
+ "add": "add", "sub": "subtract", "mul": "multiply", "truediv": "divide",
+ "floordiv": "floor_divide", "mod": "mod", "pow": "power",
+ "and": "logical_and", "or": "logical_or", "xor": "logical_xor",
+}.items():
+ setattr(DynamicArray, f"__{_dunder}__", _binop(_op))
+ setattr(DynamicArray, f"__r{_dunder}__", _binop(_op, reflected=True))
+
+for _dunder, _op in {
+ "lt": "less", "le": "less_equal", "gt": "greater", "ge": "greater_equal",
+ "eq": "equal", "ne": "not_equal",
+}.items():
+ setattr(DynamicArray, f"__{_dunder}__", _binop(_op))
+
+DynamicArray.__neg__ = _unop("negative")
+DynamicArray.__abs__ = _unop("abs")
+DynamicArray.__invert__ = _unop("logical_not")
+
+
def slice_array(array: DynamicArray, key) -> DynamicArray:
"""
Create a lazy slice of an array.
diff --git a/src/dyna_zarr/io.py b/src/dyna_zarr/io.py
index e202d46..ca194ac 100644
--- a/src/dyna_zarr/io.py
+++ b/src/dyna_zarr/io.py
@@ -19,10 +19,15 @@
import gc
from queue import Queue
-from .tiff_reader import read_tiff_lazy
+from .tiff_reader import open_tiff_zarr, read_tiff_lazy
from .codecs import Codecs
from .dynamic_array import DynamicArray
from .utils import parse_dtype
+from .operations._backend import (
+ device_context as _device_context, asnumpy as _asnumpy,
+ asnumpy_pinned as _asnumpy_pinned,
+ new_stream as _new_stream, use_stream as _use_stream,
+)
def _parse_storage_location(file_path):
@@ -36,7 +41,12 @@ def _parse_storage_location(file_path):
"""
if not isinstance(file_path, str):
file_path = str(file_path)
-
+
+ # Windows drive-letter path (e.g. C:\foo or C:/foo): urlparse would mistake
+ # the drive letter for a single-character URL scheme, so short-circuit to local.
+ if len(file_path) >= 2 and file_path[0].isalpha() and file_path[1] == ':':
+ return 'local', Path(file_path)
+
# Parse URL scheme
parsed = urlparse(file_path)
@@ -147,9 +157,9 @@ def read_file(file_path):
)
elif _is_tiff_path(str(parsed_path)):
- # Local TIFF file
- return read_tiff_lazy(parsed_path)
-
+ # Local TIFF file -> lazy zarr array (raw backend object; read_array wraps it)
+ return open_tiff_zarr(parsed_path)
+
else:
raise ValueError(
f"Unsupported file type: {parsed_path.suffix}. "
@@ -162,10 +172,8 @@ def read_file(file_path):
path_str = str(parsed_path)
if _is_tiff_path(path_str):
- # TIFF file on cloud storage
- # Note: tifffile may support remote TIFF via fsspec
- # For now, try to read via tifffile (it should handle the URL)
- return read_tiff_lazy(file_path)
+ # TIFF file on cloud storage (tifffile may support remote via fsspec)
+ return open_tiff_zarr(file_path)
else:
# Assume it's a Zarr array on cloud storage
@@ -234,7 +242,7 @@ def read_file(file_path):
if _is_tiff_path(path_str):
# TIFF over HTTP - pass to tifffile (may support via fsspec)
- return read_tiff_lazy(file_path)
+ return open_tiff_zarr(file_path)
else:
# Zarr over HTTP
spec = {
@@ -276,26 +284,10 @@ def read_array(source: Union[str, Path]) -> 'DynamicArray':
source_path = Path(source) if not isinstance(source, Path) else source
- # Check if it's a TIFF or zarr file - use read_file for both
+ # TIFF -> lazy zarr array via tifffile's bridge, wrapped as a normal zarr-backed
+ # DynamicArray (laziness/slicing/memory-bounded reads all come from DynamicArray).
if str(source_path).lower().endswith(('.tif', '.tiff')):
- # TIFF file
- ts_array = read_file(source)
- dyn_array = object.__new__(DynamicArray)
- dyn_array._ts_array = ts_array
- dyn_array._is_tensorstore = True
- dyn_array._source = source
- dyn_array._shape = tuple(ts_array.shape)
- dyn_array._dtype = ts_array.dtype
- dyn_array._transform = None
- dyn_array._chunks = ts_array.chunks if hasattr(ts_array, 'chunks') else None
- # TIFF files don't have zarr metadata, set defaults
- dyn_array._zarr_array = None
- dyn_array._zarr_format = None
- dyn_array._compressor = None
- dyn_array._compressors = None
- dyn_array._shards = None
- dyn_array._codecs = None
- return dyn_array
+ return read_tiff_lazy(source)
elif (isinstance(source_path, Path) and source_path.is_dir() and
((source_path / ".zarray").exists() or
@@ -480,6 +472,8 @@ def write_array(
gc_interval: float = 15.0,
early_quarter_timeout: Optional[float] = None,
early_tenth_timeout: Optional[float] = None,
+ max_inflight_writes: Optional[int] = None,
+ device: Optional[str] = None,
**kwargs
):
"""
@@ -518,9 +512,44 @@ def write_array(
gc_interval : float
Seconds between GC runs (default: 15.0)
"""
+ # A reshape/flatten as the OUTERMOST op is a C-order flat re-index that conflicts with
+ # nd chunk layout -- the region writer can't do it memory-bound. Route it to the
+ # streaming reindex writer (per-output-chunk, one input chunk at a time): hard-bounded
+ # to ~1 input chunk + 1 output chunk, race-free, on any chunking.
+ _tr = getattr(array, "_transform", None)
+ _trname = type(_tr).__name__
+ _local = _parse_storage_location(str(output_path))[0] == "local"
+ if _trname == "FlattenTransform" and _local:
+ # disk-staged rechunk-to-flat-contiguous + relabel: read-once source. Per-worker memory
+ # = region_size_mb; peak ~= max_workers * region_size_mb (same knobs as the region path).
+ from .rechunk import flatten_write
+ return flatten_write(_tr.array, output_path, output_chunks=chunks,
+ max_mem=int(region_size_mb * 1024 * 1024), max_workers=max_workers,
+ dtype=dtype, zarr_format=zarr_format or 2)
+ if _trname == "ReshapeTransform" and _local:
+ # staged: flatten source to 1D contiguous F (read-once, bounded), then reindex F ->
+ # target shape (contiguous input => no chunk-amplification). Per-worker memory =
+ # region_size_mb; forwards max_workers. Same knobs as the region path.
+ from .rechunk import reshape_write
+ return reshape_write(_tr.array, array.shape, output_path, output_chunks=chunks,
+ max_mem=int(region_size_mb * 1024 * 1024), max_workers=max_workers,
+ dtype=dtype, zarr_format=zarr_format or 2)
+ if _trname == "ScanTransform" and _local:
+ # bounded-carry streaming scan (tile cross-section, walk the scan axis in strips with a
+ # running carry): fully memory-bounded, read-once/write-once. Per-call budget =
+ # region_size_mb (the pull read(key) path stays prefix-bounded for lazy sub-slices).
+ from .operations.scan import scan_write
+ return scan_write(_tr.array, _tr.op, _tr.axis, output_path, output_chunks=chunks,
+ max_mem=int(region_size_mb * 1024 * 1024),
+ dtype=dtype, zarr_format=zarr_format or 2)
import tensorstore as ts
if num_readers is None:
- num_readers = max(1, max_workers * 2)
+ # One reader per writer (balanced). The async writes are the bottleneck, so extra
+ # readers just race ahead and pin more region buffers (peak RAM scales with reader
+ # count) without improving throughput -- measured 2x readers was both heavier AND
+ # slower than 1x. Raise num_readers explicitly to hide read latency on slow/remote
+ # stores. Peak RAM ~= 2 * (num_readers + queue_size + max_inflight_writes) * region_size_mb.
+ num_readers = max(1, max_workers)
# Aggressive memory cleanup before starting
gc.collect()
@@ -661,9 +690,23 @@ def write_array(
print(f"[Optimized] Total regions to process: {total_chunks}, Shape: {input_shape}, Region: {region_shape}", flush=True)
- # Queue for work distribution (good for pipeline buffering)
+ # Queue for work distribution (good for pipeline buffering).
+ # Peak RAM ~= 2 * (num_readers + queue_size + max_inflight_writes) * region_size_mb.
+ # The floor of 4 gives read-ahead/write-behind headroom so readers and writers
+ # overlap even at low worker counts. (Measured: dropping the floor to scale purely
+ # with max_workers trimmed RAM only at 1 worker, at a throughput cost, and was
+ # same-or-worse at >=2 workers -- the floor is NOT the memory driver; the staged
+ # read->queue->async-write pipeline structurally holds more live blocks than dask's
+ # single-stage model, which is also why it is faster.) Old default was
+ # min(128, max(32, num_readers)) = 32 -> a ~256MB queue ceiling regardless of size.
if queue_size is None:
- queue_size = min(128, max(32, num_readers))
+ queue_size = max(4, max_workers)
+
+ # Cap concurrent async writes: each pending tensorstore write future pins its
+ # ~region_size data buffer until it commits, so without this fast submission keeps
+ # the whole array's buffers alive at once (RSS ~= array size). See writer_thread.
+ if max_inflight_writes is None:
+ max_inflight_writes = max(4, 2 * max_workers)
chunk_queue = Queue(maxsize=queue_size)
sentinel_lock = threading.Lock()
@@ -688,6 +731,9 @@ def write_array(
def reader_thread():
"""Fast producer - no overhead."""
+ # Each reader thread owns a CUDA stream so region GPU pipelines overlap (one
+ # region's compute runs while another's H2D/D2H copies). None on CPU.
+ gpu_stream = _new_stream(device)
try:
while not shutdown_flag.is_set():
if state.get('error'): # Check for errors
@@ -714,13 +760,21 @@ def reader_thread():
for start, rs, dim_size in zip(chunk_start, region_shape, input_shape)
)
- # Read actual data using _read_direct to avoid creating SliceTransform
- data = input_array._read_direct(chunk_slice)
-
+ # Read actual data using _read_direct to avoid creating SliceTransform.
+ # device_context is thread-local, so set it here (in the reader thread)
+ # so the op chain runs on the requested device for this region; the
+ # per-thread stream lets regions overlap on the GPU. Then bring the
+ # region back to host (D2H) since tensorstore writes numpy.
+ with _use_stream(gpu_stream), _device_context(device):
+ data = input_array._read_direct(chunk_slice)
+ # pinned D2H: faster + avoids pinned-staging contention across
+ # the concurrent reader threads (measured ~7-18% faster on GPU).
+ data = _asnumpy_pinned(data)
+
# Convert dtype if needed (allows unsafe casting if explicitly requested)
if data.dtype != final_dtype_obj:
data = data.astype(final_dtype_obj, copy=False)
-
+
chunk_queue.put((chunk_slice, data))
if current_idx % 10 == 0:
@@ -776,8 +830,28 @@ def writer_thread():
# Commit futures immediately to ensure all are tracked
with futures_lock:
write_futures.append(write_future)
-
+
chunk_queue.task_done()
+
+ # Backpressure + release. Each pending write future pins its
+ # ~region_size data buffer until it commits, and a *done* future keeps
+ # pinning it until the future object itself is dropped. So we must both
+ # (a) prune completed futures here to free their buffers, and (b) block
+ # while too many writes are still in flight. Without this, buffers for
+ # the whole array stay alive at once (RSS grows with array size); with
+ # it, peak RAM stays ~ max_inflight_writes * region_size, array-independent.
+ while not shutdown_flag.is_set() and not state.get('error'):
+ with futures_lock:
+ done = [f for f in write_futures if f.done()]
+ for f in done:
+ write_futures.remove(f)
+ n_inflight = len(write_futures)
+ for f in done:
+ f.result() # surface any write error early (already complete)
+ del done
+ if n_inflight < max_inflight_writes:
+ break
+ time.sleep(0.001)
except Exception as e:
print(f"[Writer] ERROR: {e}", flush=True)
@@ -818,7 +892,12 @@ def monitor_progress():
quarter_checked = False
while not stop_monitor.is_set():
- time.sleep(2.0)
+ # Interruptible wait: Event.wait() returns immediately when stop_monitor is
+ # set, so a fast write isn't padded by a full sleep interval. Using a plain
+ # time.sleep(2.0) here left monitor.join(timeout=2.0) blocking ~2s on every
+ # write (a fixed latency floor that dominated small/medium writes).
+ if stop_monitor.wait(2.0):
+ break
try:
current_read = state['read_idx']
with futures_lock:
@@ -1005,12 +1084,15 @@ def write(
gc_interval: float = 15.0,
early_quarter_timeout: Optional[float] = None,
early_tenth_timeout: Optional[float] = None,
+ max_inflight_writes: Optional[int] = None,
+ device: Optional[str] = None,
**kwargs
):
- """Write array to Zarr. See write_array() for details."""
+ """Write array to Zarr. See write_array() for details. ``device`` ('cpu'|'cuda')
+ runs each region's op chain on that device (results are written from host)."""
return write_array(
array, output_path, max_workers, num_readers, queue_size,
chunks, shard_coefficients, dtype, compressor, zarr_format,
region_size_mb, gc_interval, early_quarter_timeout,
- early_tenth_timeout, **kwargs
+ early_tenth_timeout, max_inflight_writes, device, **kwargs
)
diff --git a/src/dyna_zarr/operations/__init__.py b/src/dyna_zarr/operations/__init__.py
new file mode 100644
index 0000000..f46a6b1
--- /dev/null
+++ b/src/dyna_zarr/operations/__init__.py
@@ -0,0 +1,39 @@
+"""operations: lazy array ops for DynamicArray, decentralized by the ome_zarr_pro
+taxonomy locality axis. Each category module owns BOTH its transforms and its public op
+functions; this package is a thin flat re-export hub. Adding an op touches exactly one
+module.
+
+ from dyna_zarr import operations as ops
+ ops.add(a, b) # flat
+ ops.gaussian_filter(x, 2.0) # flat
+ ops.neighborhood.gaussian_filter(x, 2.0) # grouped access, also available
+
+Modules by locality: structural (coordinate ops), pointwise (map_blocks + ufuncs),
+neighborhood (map_overlap + filters), reductions (min/max, streaming reduce to come).
+"""
+
+from ._base import Transform, _is_int_index, _perm_on_surviving
+
+# submodules kept importable for grouped access: operations.neighborhood.gaussian_filter, ...
+from . import structural, pointwise, reductions, neighborhood, differences, creation, scan
+
+# ...and re-exported flat: operations.gaussian_filter, operations.add, ...
+from .structural import * # noqa: F401,F403
+from .pointwise import * # noqa: F401,F403
+from .reductions import * # noqa: F401,F403
+from .neighborhood import * # noqa: F401,F403
+from .differences import * # noqa: F401,F403
+from .creation import * # noqa: F401,F403
+from .scan import * # noqa: F401,F403
+
+__all__ = (
+ ["Transform", "structural", "pointwise", "reductions", "neighborhood",
+ "differences", "creation", "scan"]
+ + structural.__all__
+ + pointwise.__all__
+ + reductions.__all__
+ + neighborhood.__all__
+ + differences.__all__
+ + creation.__all__
+ + scan.__all__
+)
diff --git a/src/dyna_zarr/operations/_backend.py b/src/dyna_zarr/operations/_backend.py
new file mode 100644
index 0000000..2351899
--- /dev/null
+++ b/src/dyna_zarr/operations/_backend.py
@@ -0,0 +1,151 @@
+"""Array-backend seam: let every transform compute with the array module of *its input
+block* rather than a hardcoded ``numpy`` / ``scipy``. On CPU the module is numpy and
+behaviour is identical; when a region's data is a CuPy array (the GPU device seam, added
+separately) the same transforms run on the GPU.
+
+cupy is detected by module name so importing this file never requires cupy to be present.
+Critically, this lets us avoid ``numpy.asarray(gpu_array)`` calls, which would silently
+copy device data back to the host mid-pipeline.
+"""
+
+import threading
+from contextlib import contextmanager
+
+import numpy as np
+
+
+def _is_cupy(x) -> bool:
+ return type(x).__module__.split(".")[0] == "cupy"
+
+
+# --------------------------------------------------------------------------- #
+# Execution-device context: the terminal (compute/io.write) sets the default
+# device that inherit-ops (device=None) run on; an op's explicit device overrides.
+# --------------------------------------------------------------------------- #
+
+_ctx = threading.local()
+
+
+def current_device() -> str:
+ return getattr(_ctx, "device", "cpu")
+
+
+@contextmanager
+def device_context(device):
+ """Set the execution device for inherit-ops within the block ('cpu' default)."""
+ prev = getattr(_ctx, "device", None)
+ _ctx.device = device or "cpu"
+ try:
+ yield
+ finally:
+ if prev is None:
+ _ctx.__dict__.pop("device", None)
+ else:
+ _ctx.device = prev
+
+
+def resolve_device(op_device) -> str:
+ """An op's explicit device if set, else the current execution context."""
+ return op_device if op_device is not None else current_device()
+
+
+def to_device(x, device):
+ """Move array ``x`` to ``device`` ('cpu' | 'cuda' | 'cuda:N'). Scalars pass through;
+ already-on-target arrays are returned untouched (no copy)."""
+ if np.isscalar(x) or x is None:
+ return x
+ dev = device or "cpu"
+ if dev == "cpu":
+ return asnumpy(x) if _is_cupy(x) else x
+ import cupy
+ if dev.startswith("cuda:"):
+ with cupy.cuda.Device(int(dev.split(":", 1)[1])):
+ return cupy.asarray(x)
+ return x if _is_cupy(x) else cupy.asarray(x)
+
+
+def array_namespace(*arrays):
+ """Return the array module (``numpy`` or ``cupy``) for the given operands. Scalars are
+ ignored; if any operand is a CuPy array, cupy is returned (and imported lazily)."""
+ for a in arrays:
+ if _is_cupy(a):
+ import cupy
+ return cupy
+ return np
+
+
+def ndimage_namespace(x):
+ """Return ``scipy.ndimage`` or ``cupyx.scipy.ndimage`` to match ``x``'s device."""
+ if _is_cupy(x):
+ import cupyx.scipy.ndimage as cndi
+ return cndi
+ import scipy.ndimage as ndi
+ return ndi
+
+
+def asnumpy(x):
+ """Host numpy array for ``x`` (D2H copy if it lives on the GPU)."""
+ if _is_cupy(x):
+ import cupy
+ return cupy.asnumpy(x)
+ return np.asarray(x)
+
+
+def asnumpy_pinned(x):
+ """Like :func:`asnumpy`, but stages the D2H copy into pinned (page-locked) host memory.
+
+ Pinned transfers are ~1.5-2x faster and, more importantly under many concurrent
+ writer threads, avoid contention on CuPy's shared pinned staging buffer (measured
+ ~7-18% faster GPU io.write on heavy ops). Used for the streamed region write path;
+ falls back to :func:`asnumpy` off-GPU or if pinned allocation is unavailable."""
+ if not _is_cupy(x):
+ return np.asarray(x)
+ try:
+ import cupyx
+ out = cupyx.empty_pinned(x.shape, dtype=x.dtype)
+ x.get(out=out) # D2H into pinned host on the current stream
+ return out
+ except Exception:
+ import cupy
+ return cupy.asnumpy(x)
+
+
+def is_gpu_array(x) -> bool:
+ return _is_cupy(x)
+
+
+def xp_for_device(device):
+ """Array module to *create* new arrays on ``device`` ('cpu'->numpy, 'cuda'->cupy).
+ Used by generative sources (creation ops) that have no input block to dispatch on."""
+ if (device or "cpu") == "cpu":
+ return np
+ import cupy
+ return cupy
+
+
+def new_stream(device):
+ """A fresh non-blocking CUDA stream for ``device`` ('cuda'/'cuda:N'), else None.
+ Giving each worker thread its own stream lets regions overlap on the GPU (one
+ region's compute runs while another's H2D/D2H copies), instead of serialising on
+ the shared default stream."""
+ dev = device or "cpu"
+ if dev == "cpu":
+ return None
+ try:
+ import cupy
+ if dev.startswith("cuda:"):
+ with cupy.cuda.Device(int(dev.split(":", 1)[1])):
+ return cupy.cuda.Stream(non_blocking=True)
+ return cupy.cuda.Stream(non_blocking=True)
+ except Exception:
+ return None
+
+
+@contextmanager
+def use_stream(stream):
+ """Activate ``stream`` as the current CUDA stream for the block (no-op if None)."""
+ if stream is None:
+ yield
+ else:
+ with stream:
+ yield
diff --git a/src/dyna_zarr/operations/_base.py b/src/dyna_zarr/operations/_base.py
new file mode 100644
index 0000000..16b162c
--- /dev/null
+++ b/src/dyna_zarr/operations/_base.py
@@ -0,0 +1,46 @@
+"""Transform base class and shared slice-math helpers for the operations package."""
+
+import numpy as np
+from typing import Tuple, Union, List, Optional, Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from dyna_zarr.dynamic_array import DynamicArray
+
+
+def _is_int_index(k) -> bool:
+ """True if key element is an integer index (drops its axis), not a slice/newaxis."""
+ return isinstance(k, (int, np.integer))
+
+
+def _perm_on_surviving(result: np.ndarray, out_key, axes) -> np.ndarray:
+ """Reorder ``result`` into output-axis order for a transpose/swapaxes read.
+
+ ``result`` was read from the underlying array with axes in *input* order, with
+ integer-indexed axes already dropped by the read. ``out_key`` is the output-space
+ key (len == len(axes)); output position ``i`` maps to input axis ``axes[i]`` and
+ ``out_key[i]`` applies to it. We transpose ``result`` into output order, restricted
+ to the axes that survive integer indexing. Reduces to ``np.transpose(result, axes)``
+ when the key contains no integer indices.
+ """
+ ndim = len(axes)
+ dropped = {axes[i] for i in range(ndim) if _is_int_index(out_key[i])}
+ surviving = [ia for ia in range(ndim) if ia not in dropped]
+ pos = {ia: j for j, ia in enumerate(surviving)}
+ desired = [axes[i] for i in range(ndim) if not _is_int_index(out_key[i])]
+ return np.transpose(result, [pos[ia] for ia in desired])
+
+
+class Transform:
+ """
+ Base class for lazy transformations.
+ """
+
+ def __init__(self):
+ self.shape = None
+ self.chunks = None
+ self.dtype = None
+
+ def read(self, key):
+ raise NotImplementedError
+
+
diff --git a/src/dyna_zarr/operations/creation.py b/src/dyna_zarr/operations/creation.py
new file mode 100644
index 0000000..ed56e2e
--- /dev/null
+++ b/src/dyna_zarr/operations/creation.py
@@ -0,0 +1,163 @@
+"""Creation ops (nullary generative sources): zeros/ones/full/empty/random (+ *_like).
+
+These build a lazy DynamicArray from a shape/dtype with no underlying array. The
+GenerativeTransform *synthesizes* whatever region a read asks for, so creation stays
+lazy and memory-bound (a huge ``zeros`` never allocates the whole thing) and composes
+with every other op. It also honours the execution device: under ``compute(device='cuda')``
+a ``zeros`` region is created directly on the GPU.
+
+``random`` is **position-deterministic**: each element's value is a hash of its GLOBAL
+index (and a per-array seed), so it is idempotent and chunk-invariant -- reading any region
+gives the same values as the whole array sliced, and ``io.write(random(...))`` equals
+``random(...).compute()``. (Unlike dask.array.random, whose values depend on the chunking.)
+"""
+
+import numpy as np
+
+from ._base import Transform, _is_int_index
+from ._backend import resolve_device, xp_for_device, to_device
+
+
+class GenerativeTransform(Transform):
+ """A nullary source: ``read(key)`` synthesizes the requested region."""
+
+ def __init__(self, shape, dtype, chunks, mode, fill_value=0, seed=None, device=None):
+ super().__init__()
+ self.shape = tuple(int(s) for s in shape)
+ self.dtype = np.dtype(dtype)
+ self.chunks = tuple(chunks) if chunks is not None else self.shape
+ self.mode = mode # 'zeros' | 'ones' | 'full' | 'empty' | 'random'
+ self.fill_value = fill_value
+ self.seed = seed
+ self._seed64 = (int(seed) & 0xFFFFFFFFFFFFFFFF) if seed is not None else 0
+ self.device = device
+
+ def _axes(self, key):
+ """Per input axis: (is_int, start, step, out_size) with concrete non-negatives."""
+ ndim = len(self.shape)
+ if not isinstance(key, tuple):
+ key = (key,)
+ key = key + (slice(None),) * (ndim - len(key))
+ axes = []
+ for a in range(ndim):
+ k = key[a]
+ size = self.shape[a]
+ if _is_int_index(k):
+ idx = int(k) if k >= 0 else size + int(k)
+ axes.append((True, idx, 1, None))
+ else:
+ start, stop, step = k.indices(size)
+ axes.append((False, start, step, len(range(start, stop, step))))
+ return axes
+
+ def _random_region(self, axes, out_shape):
+ """Position-deterministic uniform [0,1): value = splitmix64(global_flat_index, seed).
+ Built with numpy (portable, exact); moved to the device by the caller."""
+ ndim = len(self.shape)
+ strides = [1] * ndim
+ for a in range(ndim - 2, -1, -1):
+ strides[a] = strides[a + 1] * self.shape[a + 1]
+ # Work IN-PLACE in a single uint64 buffer `z` (+ one scratch `t`) so we don't
+ # allocate several full-region arrays -- important for large chunks.
+ z = np.full(out_shape, np.uint64(0x9E3779B97F4A7C15) + np.uint64(self._seed64),
+ dtype=np.uint64)
+ oi = 0
+ for a, (is_int, start, step, n) in enumerate(axes):
+ stride = np.uint64(strides[a])
+ if is_int:
+ z += np.uint64(start) * stride
+ else:
+ coord = (np.uint64(start) + np.arange(n, dtype=np.uint64) * np.uint64(step)) * stride
+ shp = [1] * len(out_shape)
+ shp[oi] = n
+ z += coord.reshape(shp) # broadcast add, in place
+ oi += 1
+ t = z >> np.uint64(30); z ^= t; z *= np.uint64(0xBF58476D1CE4E5B9)
+ np.right_shift(z, np.uint64(27), out=t); z ^= t; z *= np.uint64(0x94D049BB133111EB)
+ np.right_shift(z, np.uint64(31), out=t); z ^= t
+ del t
+ z >>= np.uint64(11)
+ u = z.astype(np.float64); u *= (1.0 / 9007199254740992.0) # /2**53 -> [0,1)
+ return u.astype(self.dtype, copy=False)
+
+ def read(self, key):
+ axes = self._axes(key)
+ out_shape = tuple(n for (is_int, _s, _st, n) in axes if not is_int)
+ dev = resolve_device(self.device)
+ if self.mode == "random":
+ return to_device(self._random_region(axes, out_shape), dev)
+ xp = xp_for_device(dev)
+ if self.mode == "zeros":
+ return xp.zeros(out_shape, dtype=self.dtype)
+ if self.mode == "ones":
+ return xp.ones(out_shape, dtype=self.dtype)
+ if self.mode == "full":
+ return xp.full(out_shape, self.fill_value, dtype=self.dtype)
+ return xp.empty(out_shape, dtype=self.dtype) # 'empty' (uninitialized)
+
+
+def _make(shape, dtype, chunks, mode, fill_value=0, seed=None, device=None):
+ from dyna_zarr.dynamic_array import DynamicArray
+ if np.isscalar(shape):
+ shape = (int(shape),)
+ t = GenerativeTransform(shape, dtype, chunks, mode, fill_value, seed, device)
+ return DynamicArray._from_transform(t)
+
+
+# --------------------------------------------------------------------------- #
+# Public ops
+# --------------------------------------------------------------------------- #
+
+def zeros(shape, dtype=np.float64, chunks=None, device=None):
+ """Lazy array of zeros."""
+ return _make(shape, dtype, chunks, "zeros", device=device)
+
+
+def ones(shape, dtype=np.float64, chunks=None, device=None):
+ """Lazy array of ones."""
+ return _make(shape, dtype, chunks, "ones", device=device)
+
+
+def full(shape, fill_value, dtype=None, chunks=None, device=None):
+ """Lazy array filled with ``fill_value``."""
+ if dtype is None:
+ dtype = np.array(fill_value).dtype
+ return _make(shape, dtype, chunks, "full", fill_value=fill_value, device=device)
+
+
+def empty(shape, dtype=np.float64, chunks=None, device=None):
+ """Lazy uninitialized array (values are undefined; each read is fresh garbage)."""
+ return _make(shape, dtype, chunks, "empty", device=device)
+
+
+def random(shape, dtype=np.float32, chunks=None, seed=None, device=None):
+ """Lazy uniform-random array in [0, 1), position-deterministic (chunk-invariant, so
+ reads are consistent and io.write matches compute). ``seed`` fixes the values; if None,
+ a fresh random seed is drawn once so separate ``random(...)`` calls differ (but each
+ array is internally consistent)."""
+ if seed is None:
+ seed = int(np.random.SeedSequence().generate_state(1, dtype=np.uint64)[0])
+ return _make(shape, dtype, chunks, "random", seed=seed, device=device)
+
+
+def zeros_like(array, dtype=None, chunks=None, device=None):
+ return zeros(array.shape, dtype=dtype or array.dtype, chunks=chunks or array.chunks, device=device)
+
+
+def ones_like(array, dtype=None, chunks=None, device=None):
+ return ones(array.shape, dtype=dtype or array.dtype, chunks=chunks or array.chunks, device=device)
+
+
+def full_like(array, fill_value, dtype=None, chunks=None, device=None):
+ return full(array.shape, fill_value, dtype=dtype or array.dtype, chunks=chunks or array.chunks, device=device)
+
+
+def empty_like(array, dtype=None, chunks=None, device=None):
+ return empty(array.shape, dtype=dtype or array.dtype, chunks=chunks or array.chunks, device=device)
+
+
+__all__ = [
+ "GenerativeTransform",
+ "zeros", "ones", "full", "empty", "random",
+ "zeros_like", "ones_like", "full_like", "empty_like",
+]
diff --git a/src/dyna_zarr/operations/differences.py b/src/dyna_zarr/operations/differences.py
new file mode 100644
index 0000000..a935bdb
--- /dev/null
+++ b/src/dyna_zarr/operations/differences.py
@@ -0,0 +1,39 @@
+"""Discrete differences: diff and gradient (halo-1 neighbourhood ops along an axis)."""
+
+from ._backend import array_namespace
+from .pointwise import subtract
+from .neighborhood import map_overlap
+
+
+def diff(array, n=1, axis=-1):
+ """Discrete difference along ``axis`` (``out[i] = a[i+1] - a[i]``), ``n`` times.
+
+ Composed from slicing + subtract, so it stays lazy and memory-bound: each read pulls
+ two adjacent strips. The axis shrinks by ``n`` (reductive, like numpy.diff).
+ """
+ ax = axis if axis >= 0 else array.ndim + axis
+ for _ in range(int(n)):
+ nd = array.ndim
+ m = array.shape[ax] # explicit positive bounds: SliceTransform
+ front = array[tuple(slice(1, m) if i == ax else slice(None) for i in range(nd))]
+ back = array[tuple(slice(0, m - 1) if i == ax else slice(None) for i in range(nd))]
+ array = subtract(front, back)
+ return array
+
+
+def gradient(array, axis=-1, device=None):
+ """Central-difference gradient along a single ``axis`` (shape-preserving, like
+ numpy.gradient with edge_order=1).
+
+ A depth-1 map_overlap: interior pixels get true central differences (the halo supplies
+ neighbours across region seams); the true array edges use ``odd_reflect`` padding
+ (linear extrapolation ``2*edge - inner``), which makes the central difference there
+ equal numpy's one-sided edge difference -- so it matches numpy exactly, everywhere.
+ """
+ ax = axis if axis >= 0 else array.ndim + axis
+ depth = tuple(1 if i == ax else 0 for i in range(array.ndim))
+ func = lambda b: array_namespace(b).gradient(b, axis=ax)
+ return map_overlap(array, func, depth, boundary="odd_reflect", name="gradient", device=device)
+
+
+__all__ = ["diff", "gradient"]
diff --git a/src/dyna_zarr/operations/neighborhood.py b/src/dyna_zarr/operations/neighborhood.py
new file mode 100644
index 0000000..8729729
--- /dev/null
+++ b/src/dyna_zarr/operations/neighborhood.py
@@ -0,0 +1,279 @@
+"""Neighbourhood primitive (pull-model ``map_overlap``) and the filters built on it.
+
+A neighbourhood op is shape-preserving but each output element depends on a local window
+of the input (halo/``depth``). In the pull model this is just map_blocks with a widened
+read window: to produce the block for ``key`` we read ``key`` expanded by ``depth`` on
+every side (padding at the true array edges per ``boundary``), apply the (shape-preserving)
+function to that expanded block, then crop the halo back off. Because every read pulls its
+own halo, the result is independent of the region size that drove it -- i.e. chunk/region
+invariant, and exact vs applying the function to the whole array (given ``depth`` >= the
+function's radius and a matching ``boundary``).
+
+This is the ``conservative . unary . neighbourhood`` taxonomy cell.
+"""
+
+import numpy as np
+from typing import Optional, Union, Tuple, Dict
+
+from ._base import Transform, _is_int_index
+from ._backend import array_namespace, ndimage_namespace, resolve_device, to_device
+
+
+# boundary name (scipy.ndimage convention) -> (numpy.pad mode, extra pad kwargs) for the
+# true array edges.
+_BOUNDARY_TO_NPPAD = {
+ "reflect": ("symmetric", {}), # scipy 'reflect': (d c b a | a b c d) -- edge duplicated
+ "mirror": ("reflect", {}), # scipy 'mirror': (d c b | a b c d) -- edge not duplicated
+ "nearest": ("edge", {}),
+ "wrap": ("wrap", {}),
+ "constant": ("constant", {}),
+ # odd-reflection = linear extrapolation (2*edge - inner); makes a central difference at
+ # the boundary equal the one-sided difference, so gradient() matches numpy at the edges.
+ "odd_reflect": ("reflect", {"reflect_type": "odd"}),
+}
+
+
+def _normalize_depth(depth, ndim) -> Tuple[int, ...]:
+ """Normalize ``depth`` to a per-axis tuple of symmetric halo widths."""
+ if isinstance(depth, dict):
+ return tuple(int(depth.get(a, 0)) for a in range(ndim))
+ if np.isscalar(depth):
+ return (int(depth),) * ndim
+ depth = tuple(int(d) for d in depth)
+ if len(depth) != ndim:
+ raise ValueError(f"depth {depth} does not match array ndim {ndim}")
+ return depth
+
+
+def _infer_overlap_dtype(func, array):
+ """Infer output dtype by applying func to a small sample of the array's dtype."""
+ from ..utils import parse_dtype
+ np_dt = parse_dtype(array.dtype)[0]
+ sample_shape = tuple(min(s, 3) for s in array.shape)
+ try:
+ return np.asarray(func(np.ones(sample_shape, dtype=np_dt))).dtype
+ except Exception:
+ return np_dt
+
+
+class MapOverlapTransform(Transform):
+ """Apply a shape-preserving neighbourhood ``func`` with a ``depth`` halo, lazily.
+
+ ``func`` receives an expanded block (input window padded to include the halo) and must
+ return an array of the *same shape*; the halo is then trimmed. ``depth`` may be an int
+ (all axes), a per-axis sequence, or a ``{axis: depth}`` dict. ``boundary`` (scipy.ndimage
+ names: reflect/mirror/nearest/wrap/constant) controls how the true array edges are
+ extended. For an exact match to ``func`` applied to the whole array, ``depth`` must be
+ >= the function's radius and ``boundary`` must match the function's edge mode.
+ """
+
+ def __init__(self, array, func, depth, boundary="reflect", dtype=None, name=None,
+ device=None):
+ super().__init__()
+ if boundary not in _BOUNDARY_TO_NPPAD:
+ raise ValueError(f"unknown boundary {boundary!r}; expected one of "
+ f"{sorted(_BOUNDARY_TO_NPPAD)}")
+ self.array = array
+ self.func = func
+ self.device = device
+ self.name = name or getattr(func, "__name__", "map_overlap")
+ self.depth = _normalize_depth(depth, array.ndim)
+ self.boundary = boundary
+ self.shape = array.shape
+ self.chunks = array.chunks
+ self.dtype = np.dtype(dtype) if dtype is not None else _infer_overlap_dtype(func, array)
+
+ def read(self, key):
+ ndim = self.array.ndim
+ if not isinstance(key, tuple):
+ key = (key,)
+ key = key + (slice(None),) * (ndim - len(key))
+
+ read_slices = [] # region to read from the input (clamped to bounds, step 1)
+ pad_widths = [] # (before, after) padding that restores the halo at true edges
+ crop_slices = [] # crop the func output back to the core region + apply step
+ squeeze_axes = []
+ for a in range(ndim):
+ k = key[a]
+ size = self.array.shape[a]
+ d = self.depth[a]
+ if _is_int_index(k):
+ idx = int(k) if k >= 0 else size + int(k)
+ start, stop, step = idx, idx + 1, 1
+ squeeze_axes.append(a)
+ else:
+ start, stop, step = k.indices(size)
+ if step < 0:
+ raise NotImplementedError("map_overlap: negative-step reads not supported")
+
+ # desired expanded span [start-d, stop+d), clamped to [0, size)
+ read_slices.append(slice(max(0, start - d), min(size, stop + d)))
+ pad_widths.append((max(0, d - start), max(0, (stop + d) - size)))
+ # after padding, the block has the core at [d : d+core_len]; then apply the step
+ crop_slices.append(slice(d, d + (stop - start), step))
+
+ block = to_device(self.array._read_direct(tuple(read_slices)),
+ resolve_device(self.device))
+ xp = array_namespace(block)
+ if any(pb or pa for pb, pa in pad_widths):
+ mode, pad_kw = _BOUNDARY_TO_NPPAD[self.boundary]
+ block = xp.pad(block, pad_widths, mode=mode, **pad_kw)
+
+ out = self.func(block) # func dispatches ndimage on block's device
+ out = out[tuple(crop_slices)]
+ for a in sorted(squeeze_axes, reverse=True):
+ out = xp.squeeze(out, axis=a)
+ return out
+
+
+# --------------------------------------------------------------------------- #
+# depth helpers for the scipy-backed filters
+# --------------------------------------------------------------------------- #
+
+def _as_per_axis(value, ndim):
+ if np.isscalar(value):
+ return (value,) * ndim
+ value = tuple(value)
+ if len(value) != ndim:
+ raise ValueError(f"expected a scalar or length-{ndim} sequence, got {value}")
+ return value
+
+
+def _gaussian_depth(sigma, ndim, truncate):
+ sig = _as_per_axis(sigma, ndim)
+ # scipy's 1-D Gaussian half-width: int(truncate * sigma + 0.5)
+ return tuple(int(truncate * float(s) + 0.5) for s in sig)
+
+
+def _size_depth(size, ndim):
+ sz = _as_per_axis(size, ndim)
+ # max half-extent of a length-n window (covers even sizes too)
+ return tuple(int(n) // 2 for n in sz)
+
+
+# --------------------------------------------------------------------------- #
+# Public ops -- map_overlap primitive + scipy-backed neighbourhood filters
+# --------------------------------------------------------------------------- #
+
+def map_overlap(array, func, depth, boundary="reflect", dtype=None, name=None, device=None):
+ """Apply a shape-preserving neighbourhood ``func`` with a ``depth`` halo, lazily and
+ chunk-invariantly. Every read pulls its own halo, so the result is independent of the
+ region size and exact vs applying ``func`` to the whole array (given depth >= radius and
+ matching ``boundary``). The filters below wrap this. ``device`` (None=inherit, 'cpu',
+ 'cuda') runs it on that device."""
+ return array._with_transform(
+ MapOverlapTransform(array, func, depth, boundary=boundary, dtype=dtype,
+ name=name, device=device)
+ )
+
+
+def gaussian_filter(array, sigma, boundary="reflect", truncate=4.0, device=None, **kw):
+ depth = _gaussian_depth(sigma, array.ndim, truncate)
+ func = lambda b: ndimage_namespace(b).gaussian_filter(
+ b, sigma=sigma, mode=boundary, truncate=truncate, **kw)
+ return map_overlap(array, func, depth, boundary=boundary, name="gaussian_filter", device=device)
+
+
+def uniform_filter(array, size, boundary="reflect", device=None, **kw):
+ func = lambda b: ndimage_namespace(b).uniform_filter(b, size=size, mode=boundary, **kw)
+ return map_overlap(array, func, _size_depth(size, array.ndim), boundary=boundary,
+ name="uniform_filter", device=device)
+
+
+def median_filter(array, size, boundary="reflect", device=None, **kw):
+ func = lambda b: ndimage_namespace(b).median_filter(b, size=size, mode=boundary, **kw)
+ return map_overlap(array, func, _size_depth(size, array.ndim), boundary=boundary,
+ name="median_filter", device=device)
+
+
+def minimum_filter(array, size, boundary="reflect", device=None, **kw):
+ func = lambda b: ndimage_namespace(b).minimum_filter(b, size=size, mode=boundary, **kw)
+ return map_overlap(array, func, _size_depth(size, array.ndim), boundary=boundary,
+ name="minimum_filter", device=device)
+
+
+def maximum_filter(array, size, boundary="reflect", device=None, **kw):
+ func = lambda b: ndimage_namespace(b).maximum_filter(b, size=size, mode=boundary, **kw)
+ return map_overlap(array, func, _size_depth(size, array.ndim), boundary=boundary,
+ name="maximum_filter", device=device)
+
+
+def grey_erosion(array, size, boundary="reflect", device=None, **kw):
+ func = lambda b: ndimage_namespace(b).grey_erosion(b, size=size, mode=boundary, **kw)
+ return map_overlap(array, func, _size_depth(size, array.ndim), boundary=boundary,
+ name="grey_erosion", device=device)
+
+
+def grey_dilation(array, size, boundary="reflect", device=None, **kw):
+ func = lambda b: ndimage_namespace(b).grey_dilation(b, size=size, mode=boundary, **kw)
+ return map_overlap(array, func, _size_depth(size, array.ndim), boundary=boundary,
+ name="grey_dilation", device=device)
+
+
+def _kernel_depth(weights, ndim):
+ """Halo per axis for a convolution/correlation kernel: half the kernel size (>= the
+ kernel's reach on either side, so map_overlap stays exact vs the whole-array result)."""
+ w = np.asarray(weights)
+ if w.ndim != ndim:
+ raise ValueError(
+ f"weights ndim {w.ndim} must match array ndim {ndim}; use size-1 axes to leave "
+ f"an axis untouched (e.g. a (1, ky, kx) kernel over a (z, y, x) volume)"
+ )
+ return w, tuple(int(s) // 2 for s in w.shape)
+
+
+def convolve(array, weights, boundary="reflect", cval=0.0, device=None):
+ """Multidimensional convolution with a ``weights`` kernel (like scipy.ndimage.convolve).
+ ``weights`` must have the same ndim as ``array`` (size-1 axes leave an axis untouched).
+ The kernel is moved to the block's namespace, so this runs on CPU or GPU per ``device``."""
+ w, depth = _kernel_depth(weights, array.ndim)
+ func = lambda b: ndimage_namespace(b).convolve(
+ b, array_namespace(b).asarray(w), mode=boundary, cval=cval)
+ return map_overlap(array, func, depth, boundary=boundary, name="convolve", device=device)
+
+
+def correlate(array, weights, boundary="reflect", cval=0.0, device=None):
+ """Multidimensional cross-correlation with a ``weights`` kernel (like
+ scipy.ndimage.correlate). Same as ``convolve`` but the kernel is not flipped."""
+ w, depth = _kernel_depth(weights, array.ndim)
+ func = lambda b: ndimage_namespace(b).correlate(
+ b, array_namespace(b).asarray(w), mode=boundary, cval=cval)
+ return map_overlap(array, func, depth, boundary=boundary, name="correlate", device=device)
+
+
+def laplace(array, boundary="reflect", device=None, **kw):
+ """Laplace filter via the second-derivative [1, -2, 1] stencil (scipy.ndimage.laplace).
+ Fixed radius 1 on every axis."""
+ func = lambda b: ndimage_namespace(b).laplace(b, mode=boundary, **kw)
+ return map_overlap(array, func, (1,) * array.ndim, boundary=boundary,
+ name="laplace", device=device)
+
+
+def gaussian_laplace(array, sigma, boundary="reflect", truncate=4.0, device=None, **kw):
+ """Laplace of Gaussian (scipy.ndimage.gaussian_laplace). Halo is the Gaussian radius
+ ``int(truncate * sigma + 0.5)`` per axis, as for ``gaussian_filter``."""
+ depth = _gaussian_depth(sigma, array.ndim, truncate)
+ func = lambda b: ndimage_namespace(b).gaussian_laplace(
+ b, sigma=sigma, mode=boundary, truncate=truncate, **kw)
+ return map_overlap(array, func, depth, boundary=boundary,
+ name="gaussian_laplace", device=device)
+
+
+def gaussian_gradient_magnitude(array, sigma, boundary="reflect", truncate=4.0,
+ device=None, **kw):
+ """Gradient magnitude using Gaussian derivatives (scipy.ndimage.
+ gaussian_gradient_magnitude). Halo is the Gaussian radius per axis."""
+ depth = _gaussian_depth(sigma, array.ndim, truncate)
+ func = lambda b: ndimage_namespace(b).gaussian_gradient_magnitude(
+ b, sigma=sigma, mode=boundary, truncate=truncate, **kw)
+ return map_overlap(array, func, depth, boundary=boundary,
+ name="gaussian_gradient_magnitude", device=device)
+
+
+__all__ = [
+ "MapOverlapTransform", "map_overlap",
+ "gaussian_filter", "uniform_filter", "median_filter",
+ "minimum_filter", "maximum_filter", "grey_erosion", "grey_dilation",
+ "convolve", "correlate",
+ "laplace", "gaussian_laplace", "gaussian_gradient_magnitude",
+]
diff --git a/src/dyna_zarr/operations/pointwise.py b/src/dyna_zarr/operations/pointwise.py
new file mode 100644
index 0000000..4b201d7
--- /dev/null
+++ b/src/dyna_zarr/operations/pointwise.py
@@ -0,0 +1,173 @@
+"""Pull-model map_blocks primitive: elementwise (chunk-independent) ops, memory-bound."""
+
+import numpy as np
+from typing import Tuple, Union, List, Optional, Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from dyna_zarr.dynamic_array import DynamicArray
+
+
+from ._base import Transform
+from ._backend import array_namespace, resolve_device, to_device
+
+
+def _infer_mapblocks_dtype(func, operands):
+ """Infer an elementwise func's output dtype by applying it to 1-element samples of
+ each operand's dtype (scalars passed through), letting NumPy type promotion decide.
+
+ Operand dtypes may be tensorstore dtypes (io.read yields TS-backed arrays), which
+ numpy can't consume directly, so normalise via parse_dtype first."""
+ from dyna_zarr.dynamic_array import DynamicArray
+ from ..utils import parse_dtype
+ samples = [
+ np.ones((1,), dtype=parse_dtype(o.dtype)[0]) if isinstance(o, DynamicArray) else o
+ for o in operands
+ ]
+ return np.asarray(func(*samples)).dtype
+
+
+class MapBlocksTransform(Transform):
+ """Pull-model ``map_blocks``: apply an ELEMENTWISE (chunk-independent) function to
+ one or more equally-shaped DynamicArray operands (plus scalars).
+
+ Because the function is pointwise, ``func(a[key], b[key]) == func(a, b)[key]``, so
+ reading an arbitrary output slice just reads that slice from each operand and applies
+ the function -- nothing is materialised, so it stays memory-bound. This is the
+ ``conservative . pointwise`` taxonomy cell; it is NOT valid for neighbourhood or
+ reducing functions (use map_overlap / reduce for those).
+
+ Operands may mix DynamicArrays (read blockwise) and plain scalars (passed through;
+ NumPy broadcasts them). All DynamicArray operands must share the same shape.
+ """
+
+ def __init__(self, func, operands, dtype=None, name=None, device=None):
+ super().__init__()
+ from dyna_zarr.dynamic_array import DynamicArray
+ self.func = func
+ self.operands = list(operands)
+ self.device = device
+ self.name = name or getattr(func, "__name__", "map_blocks")
+ arrays = [o for o in self.operands if isinstance(o, DynamicArray)]
+ if not arrays:
+ raise ValueError("map_blocks requires at least one DynamicArray operand")
+ ref = arrays[0]
+ for a in arrays[1:]:
+ if a.shape != ref.shape:
+ raise ValueError(
+ f"map_blocks operands must share shape; got {a.shape} vs {ref.shape}"
+ )
+ self.shape = ref.shape
+ self.chunks = ref.chunks
+ self.dtype = dtype if dtype is not None else _infer_mapblocks_dtype(func, self.operands)
+
+ def read(self, key):
+ from dyna_zarr.dynamic_array import DynamicArray
+ dev = resolve_device(self.device)
+ blocks = [
+ to_device(o._read_direct(key), dev) if isinstance(o, DynamicArray)
+ else o
+ for o in self.operands
+ ]
+ # xp.asarray keeps the result on its own device (numpy or cupy); a plain
+ # np.asarray would force a GPU->host copy mid-pipeline.
+ xp = array_namespace(*blocks)
+ return xp.asarray(self.func(*blocks))
+
+
+# --------------------------------------------------------------------------- #
+# Public ops -- pointwise ufuncs, all thin wrappers over map_blocks
+# --------------------------------------------------------------------------- #
+
+def map_blocks(func, *operands, dtype=None, name=None, device=None):
+ """Apply an elementwise (chunk-independent) ``func`` to one or more equally-shaped
+ DynamicArrays (plus scalars), lazily and memory-bound. Every pointwise op here (ufuncs,
+ comparisons, logical ops, where, clip, astype) is a thin wrapper over this. ``func``
+ receives the read block of each operand; for neighbourhood/reducing funcs use
+ map_overlap / reduce. ``device`` (None=inherit the execution context, 'cpu', 'cuda')
+ runs this op on that device."""
+ from dyna_zarr.dynamic_array import DynamicArray
+ transform = MapBlocksTransform(func, operands, dtype=dtype, name=name, device=device)
+ ref = next(o for o in operands if isinstance(o, DynamicArray))
+ return ref._with_transform(transform)
+
+
+# unary
+def abs(array): return map_blocks(np.abs, array, name="abs")
+def negative(array): return map_blocks(np.negative, array, name="negative")
+def sign(array): return map_blocks(np.sign, array, name="sign")
+def sqrt(array): return map_blocks(np.sqrt, array, name="sqrt")
+def square(array): return map_blocks(np.square, array, name="square")
+def exp(array): return map_blocks(np.exp, array, name="exp")
+def log(array): return map_blocks(np.log, array, name="log")
+def log2(array): return map_blocks(np.log2, array, name="log2")
+def log10(array): return map_blocks(np.log10, array, name="log10")
+def floor(array): return map_blocks(np.floor, array, name="floor")
+def ceil(array): return map_blocks(np.ceil, array, name="ceil")
+def reciprocal(array): return map_blocks(np.reciprocal, array, name="reciprocal")
+
+
+def round(array, decimals=0):
+ return map_blocks(lambda x: np.round(x, decimals=decimals), array, name="round")
+
+
+def clip(array, a_min=None, a_max=None):
+ return map_blocks(lambda x: np.clip(x, a_min, a_max), array, name="clip")
+
+
+def astype(array, dtype):
+ return map_blocks(lambda x: x.astype(dtype), array, dtype=np.dtype(dtype), name="astype")
+
+
+# binary
+def add(a, b): return map_blocks(np.add, a, b, name="add")
+def subtract(a, b): return map_blocks(np.subtract, a, b, name="subtract")
+def multiply(a, b): return map_blocks(np.multiply, a, b, name="multiply")
+def divide(a, b): return map_blocks(np.divide, a, b, name="divide")
+def floor_divide(a, b): return map_blocks(np.floor_divide, a, b, name="floor_divide")
+def mod(a, b): return map_blocks(np.mod, a, b, name="mod")
+def power(a, b): return map_blocks(np.power, a, b, name="power")
+def maximum(a, b): return map_blocks(np.maximum, a, b, name="maximum")
+def minimum(a, b): return map_blocks(np.minimum, a, b, name="minimum")
+
+
+# comparisons & logical (bool out)
+def greater(a, b): return map_blocks(np.greater, a, b, name="greater")
+def greater_equal(a, b): return map_blocks(np.greater_equal, a, b, name="greater_equal")
+def less(a, b): return map_blocks(np.less, a, b, name="less")
+def less_equal(a, b): return map_blocks(np.less_equal, a, b, name="less_equal")
+def equal(a, b): return map_blocks(np.equal, a, b, name="equal")
+def not_equal(a, b): return map_blocks(np.not_equal, a, b, name="not_equal")
+def logical_and(a, b): return map_blocks(np.logical_and, a, b, name="logical_and")
+def logical_or(a, b): return map_blocks(np.logical_or, a, b, name="logical_or")
+def logical_xor(a, b): return map_blocks(np.logical_xor, a, b, name="logical_xor")
+def logical_not(array): return map_blocks(np.logical_not, array, name="logical_not")
+
+
+# ternary
+def where(condition, x, y): return map_blocks(np.where, condition, x, y, name="where")
+
+
+# lookup / binning (pointwise; func resolves the array module of the block so it runs on
+# whichever device the block lives on)
+def isin(array, test_elements, invert=False):
+ """Elementwise membership test against ``test_elements`` (bool output)."""
+ return map_blocks(lambda x: array_namespace(x).isin(x, test_elements, invert=invert),
+ array, dtype=bool, name="isin")
+
+
+def digitize(array, bins, right=False):
+ """Index of the bin each element falls into (like numpy.digitize)."""
+ return map_blocks(lambda x: array_namespace(x).digitize(x, bins, right=right),
+ array, dtype=np.intp, name="digitize")
+
+
+__all__ = [
+ "MapBlocksTransform", "map_blocks",
+ "abs", "negative", "sign", "sqrt", "square", "exp", "log", "log2", "log10",
+ "floor", "ceil", "reciprocal", "round", "clip", "astype",
+ "add", "subtract", "multiply", "divide", "floor_divide", "mod", "power",
+ "maximum", "minimum",
+ "greater", "greater_equal", "less", "less_equal", "equal", "not_equal",
+ "logical_and", "logical_or", "logical_xor", "logical_not", "where",
+ "isin", "digitize",
+]
diff --git a/src/dyna_zarr/operations/reductions.py b/src/dyna_zarr/operations/reductions.py
new file mode 100644
index 0000000..aeb084e
--- /dev/null
+++ b/src/dyna_zarr/operations/reductions.py
@@ -0,0 +1,381 @@
+"""Reductions -- the reductive/destructive taxonomy cell.
+
+A reduction collapses one or more axes. The pull model can't cheaply slice a *reduced*
+axis (the output element depends on the whole axis), so ``read(key)`` reads only the
+kept-axis tile the key asks for but **streams the reduced axis in bounded chunks**,
+combining them with an associative reducer (min/max/sum/prod/mean/any/all). That keeps it
+memory-bound even for a huge reduced axis, and the result is chunk/region-invariant and
+exact vs numpy. Step/integer indices on the output are applied after the reduce (same
+crop-at-the-end trick as map_overlap).
+
+``axis=None`` reduces everything to a 0-d result, still streamed. The result is a lazy
+DynamicArray (chainable / writable); DynamicArray.min()/max() wrap it and compute eagerly.
+"""
+
+import builtins
+import numpy as np
+from typing import Optional, Tuple
+
+from ._base import Transform, _is_int_index
+from ._backend import array_namespace, asnumpy, resolve_device, to_device
+
+_DEFAULT_STRIP_BYTES = 64 * 1024 * 1024 # per-read streaming budget for the reduced axis
+
+
+def _ident(xp, s, count, ddof):
+ return s
+
+
+def _div_count(xp, s, count, ddof):
+ return s / count
+
+
+def _var_partial(xp, b, ax):
+ # accumulate in float64 so the one-pass sum-of-squares identity stays accurate
+ bf = b.astype(xp.float64)
+ return (xp.sum(bf, axis=ax), xp.sum(bf * bf, axis=ax))
+
+
+def _pair_add(xp, a, b):
+ return (a[0] + b[0], a[1] + b[1])
+
+
+def _var_finalize(xp, s, count, ddof):
+ ssum, ssq = s
+ var = (ssq - ssum * ssum / count) / (count - ddof)
+ return xp.clip(var, 0.0, None) # guard tiny negatives from float error
+
+
+def _std_finalize(xp, s, count, ddof):
+ return xp.sqrt(_var_finalize(xp, s, count, ddof))
+
+
+class _Reducer:
+ """Associative reducers use partial(xp, block, axes)->state, combine(xp, a, b)->state,
+ finalize(xp, state, count, ddof)->result and are streamed over the reduced axis.
+ Non-associative reducers (argmin/argmax) can't be chunked on the reduced axis, so they
+ read it whole per kept-tile and apply direct(xp, block, axes) in one shot."""
+ def __init__(self, partial=None, combine=None, finalize=_ident,
+ associative=True, direct=None):
+ self.partial = partial
+ self.combine = combine
+ self.finalize = finalize
+ self.associative = associative
+ self.direct = direct
+
+
+def _arg_axis(ax):
+ if len(ax) == 1:
+ return ax[0]
+ if len(ax) == 0:
+ return None
+ raise ValueError("argmin/argmax take a single axis or axis=None")
+
+
+_REDUCERS = {
+ "min": _Reducer(lambda xp, b, ax: xp.min(b, axis=ax), lambda xp, a, b: xp.minimum(a, b)),
+ "max": _Reducer(lambda xp, b, ax: xp.max(b, axis=ax), lambda xp, a, b: xp.maximum(a, b)),
+ "sum": _Reducer(lambda xp, b, ax: xp.sum(b, axis=ax), lambda xp, a, b: xp.add(a, b)),
+ "prod": _Reducer(lambda xp, b, ax: xp.prod(b, axis=ax), lambda xp, a, b: xp.multiply(a, b)),
+ "any": _Reducer(lambda xp, b, ax: xp.any(b, axis=ax), lambda xp, a, b: xp.logical_or(a, b)),
+ "all": _Reducer(lambda xp, b, ax: xp.all(b, axis=ax), lambda xp, a, b: xp.logical_and(a, b)),
+ "mean": _Reducer(lambda xp, b, ax: xp.sum(b, axis=ax), lambda xp, a, b: xp.add(a, b), _div_count),
+ "var": _Reducer(_var_partial, _pair_add, _var_finalize),
+ "std": _Reducer(_var_partial, _pair_add, _std_finalize),
+ # non-associative: read the reduced axis whole per kept-tile, apply in one shot
+ "argmin": _Reducer(associative=False, direct=lambda xp, b, ax: xp.argmin(b, axis=_arg_axis(ax))),
+ "argmax": _Reducer(associative=False, direct=lambda xp, b, ax: xp.argmax(b, axis=_arg_axis(ax))),
+ "median": _Reducer(associative=False, direct=lambda xp, b, ax: xp.median(b, axis=ax)),
+}
+
+
+def _normalize_axes(axis, ndim) -> Tuple[int, ...]:
+ if axis is None:
+ return tuple(range(ndim))
+ if np.isscalar(axis):
+ axis = (axis,)
+ out = []
+ for a in axis:
+ a = int(a) if a >= 0 else ndim + int(a)
+ if a < 0 or a >= ndim:
+ raise ValueError(f"axis {a} out of bounds for ndim {ndim}")
+ out.append(a)
+ return tuple(sorted(set(out)))
+
+
+class ReduceTransform(Transform):
+ """Lazy, streaming reduction over one or more axes."""
+
+ def __init__(self, array, reducer, axis=None, keepdims=False,
+ strip_bytes=_DEFAULT_STRIP_BYTES, device=None, ddof=0):
+ super().__init__()
+ if reducer not in _REDUCERS:
+ raise ValueError(f"unknown reducer {reducer!r}; expected {sorted(_REDUCERS)}")
+ from ..utils import parse_dtype
+ self.array = array
+ self.reducer_name = reducer
+ self.reducer = _REDUCERS[reducer]
+ self.keepdims = keepdims
+ self.strip_bytes = strip_bytes
+ self.device = device
+ self.ddof = ddof
+ ndim = array.ndim
+ self.R = _normalize_axes(axis, ndim)
+ self.kept = tuple(a for a in range(ndim) if a not in self.R)
+ self._in_dtype = parse_dtype(array.dtype)[0]
+
+ chunks = array.chunks
+ if keepdims:
+ self.shape = tuple(1 if a in self.R else array.shape[a] for a in range(ndim))
+ self.chunks = (tuple(1 if a in self.R else chunks[a] for a in range(ndim))
+ if chunks else None)
+ else:
+ self.shape = tuple(array.shape[a] for a in self.kept)
+ self.chunks = tuple(chunks[a] for a in self.kept) if chunks else None
+
+ self.dtype = self._infer_dtype()
+
+ def _infer_dtype(self):
+ sample = np.ones((2,) * self.array.ndim, dtype=self._in_dtype)
+ if not self.reducer.associative:
+ res = self.reducer.direct(np, sample, self.R)
+ else:
+ state = self.reducer.partial(np, sample, self.R)
+ res = self.reducer.finalize(np, state, 2 ** len(self.R), self.ddof)
+ return np.asarray(res).dtype
+
+ def _stream(self, input_slices):
+ """Reduce over self.R. Associative reducers stream the largest reduced axis in
+ bounded chunks + combine; non-associative ones read the reduced axis whole."""
+ arr = self.array
+ if not self.reducer.associative:
+ block = to_device(arr._read_direct(tuple(input_slices)), resolve_device(self.device))
+ xp = array_namespace(block)
+ if block.size == 0: # empty kept region: numpy median/argmin choke -> build empty
+ empty_shape = tuple(s for i, s in enumerate(block.shape) if i not in self.R)
+ return xp.empty(empty_shape, dtype=self.dtype)
+ return self.reducer.direct(xp, block, self.R)
+ kept_elems = 1
+ for a in self.kept:
+ s = input_slices[a]
+ kept_elems *= (s.stop - s.start)
+ chunk_axis = builtins.max(self.R, key=lambda a: arr.shape[a])
+ other_reduced = 1
+ for a in self.R:
+ if a != chunk_axis:
+ other_reduced *= arr.shape[a]
+ denom = builtins.max(1, kept_elems * other_reduced * self._in_dtype.itemsize)
+ chunk_len = builtins.max(1, int(self.strip_bytes // denom))
+
+ size = arr.shape[chunk_axis]
+ acc = None
+ for c in range(0, size, chunk_len):
+ slices = list(input_slices)
+ slices[chunk_axis] = slice(c, builtins.min(size, c + chunk_len))
+ block = to_device(arr._read_direct(tuple(slices)), resolve_device(self.device))
+ xp = array_namespace(block)
+ p = self.reducer.partial(xp, block, self.R)
+ acc = p if acc is None else self.reducer.combine(xp, acc, p)
+ count = 1
+ for a in self.R:
+ count *= arr.shape[a]
+ return self.reducer.finalize(xp, acc, count, self.ddof)
+
+ def read(self, key):
+ out_ndim = len(self.shape)
+ if not isinstance(key, tuple):
+ key = (key,)
+ key = key + (slice(None),) * (out_ndim - len(key))
+
+ input_slices = [slice(None)] * self.array.ndim # reduced axes stay full (streamed)
+ params = [] # (start, stop, step, is_int) per out axis
+ for oi in range(out_ndim):
+ k = key[oi]
+ osize = self.shape[oi]
+ if _is_int_index(k):
+ idx = int(k) if k >= 0 else osize + int(k)
+ start, stop, step, is_int = idx, idx + 1, 1, True
+ else:
+ start, stop, step = k.indices(osize)
+ is_int = False
+ params.append((start, stop, step, is_int))
+ in_axis = oi if self.keepdims else self.kept[oi]
+ if in_axis in self.kept:
+ # contiguous read on this kept input axis; step/int applied after reduce
+ input_slices[in_axis] = slice(start, stop)
+ # keepdims reduced axis: read stays full; its size-1 output handled in the crop
+
+ block = self._stream(input_slices) # kept-axis order, reduced axes gone
+ xp = array_namespace(block)
+ if self.keepdims:
+ for a in sorted(self.R):
+ block = xp.expand_dims(block, a)
+
+ crop = tuple(slice(0, stop - start, step) for (start, stop, step, _) in params)
+ block = block[crop]
+ for oi in sorted((i for i, p in enumerate(params) if p[3]), reverse=True):
+ block = xp.squeeze(block, axis=oi)
+ return block
+
+
+# --------------------------------------------------------------------------- #
+# Public ops -- reductions
+# --------------------------------------------------------------------------- #
+
+def reduce(array, reducer, axis=None, keepdims=False, device=None, ddof=0):
+ """Lazy streaming reduction with a named ``reducer`` (min/max/sum/prod/mean/var/std/
+ any/all/argmin/argmax) over ``axis`` (int, tuple, or None for all). ``device``
+ (None=inherit, 'cpu', 'cuda') runs it on that device; ``ddof`` applies to var/std."""
+ return array._with_transform(
+ ReduceTransform(array, reducer, axis=axis, keepdims=keepdims, device=device, ddof=ddof))
+
+
+def min(array, axis=None, keepdims=False, device=None):
+ """Minimum along axis/axes (lazy, streaming reduction)."""
+ return reduce(array, "min", axis=axis, keepdims=keepdims, device=device)
+
+
+def max(array, axis=None, keepdims=False, device=None):
+ """Maximum along axis/axes (lazy, streaming reduction)."""
+ return reduce(array, "max", axis=axis, keepdims=keepdims, device=device)
+
+
+def sum(array, axis=None, keepdims=False, device=None):
+ """Sum along axis/axes (lazy, streaming reduction)."""
+ return reduce(array, "sum", axis=axis, keepdims=keepdims, device=device)
+
+
+def prod(array, axis=None, keepdims=False, device=None):
+ """Product along axis/axes (lazy, streaming reduction)."""
+ return reduce(array, "prod", axis=axis, keepdims=keepdims, device=device)
+
+
+def mean(array, axis=None, keepdims=False, device=None):
+ """Mean along axis/axes (lazy, streaming: running sum / count)."""
+ return reduce(array, "mean", axis=axis, keepdims=keepdims, device=device)
+
+
+def any(array, axis=None, keepdims=False, device=None):
+ """Logical OR along axis/axes (lazy, streaming reduction)."""
+ return reduce(array, "any", axis=axis, keepdims=keepdims, device=device)
+
+
+def all(array, axis=None, keepdims=False, device=None):
+ """Logical AND along axis/axes (lazy, streaming reduction)."""
+ return reduce(array, "all", axis=axis, keepdims=keepdims, device=device)
+
+
+def var(array, axis=None, keepdims=False, ddof=0, device=None):
+ """Variance along axis/axes (lazy, streaming: (sum, sum-of-squares, count), float64)."""
+ return reduce(array, "var", axis=axis, keepdims=keepdims, device=device, ddof=ddof)
+
+
+def std(array, axis=None, keepdims=False, ddof=0, device=None):
+ """Standard deviation along axis/axes (lazy, streaming; sqrt of var)."""
+ return reduce(array, "std", axis=axis, keepdims=keepdims, device=device, ddof=ddof)
+
+
+def argmin(array, axis=None, keepdims=False, device=None):
+ """Index of the minimum along a single ``axis`` (or flat if None). Reads the reduced
+ axis whole per kept-tile (not associatively streamable)."""
+ return reduce(array, "argmin", axis=axis, keepdims=keepdims, device=device)
+
+
+def argmax(array, axis=None, keepdims=False, device=None):
+ """Index of the maximum along a single ``axis`` (or flat if None). Reads the reduced
+ axis whole per kept-tile (not associatively streamable)."""
+ return reduce(array, "argmax", axis=axis, keepdims=keepdims, device=device)
+
+
+def median(array, axis=None, keepdims=False, device=None):
+ """Median along axis/axes (or all if None). Not associatively streamable, so it reads
+ the reduced axis/axes whole per kept-tile (memory bounded by the kept region)."""
+ return reduce(array, "median", axis=axis, keepdims=keepdims, device=device)
+
+
+def _iter_region_slices(shape, itemsize, budget):
+ """Yield memory-bounded region slice-tuples covering ``shape`` (<= ``budget`` bytes
+ each), expanding trailing (contiguous) axes first."""
+ import itertools
+ budget_elems = builtins.max(1, int(budget) // builtins.max(1, itemsize))
+ region = [1] * len(shape)
+ acc = 1
+ for a in reversed(range(len(shape))):
+ region[a] = builtins.min(shape[a], builtins.max(1, budget_elems // acc))
+ acc *= region[a]
+ if region[a] < shape[a]:
+ break
+ for start in itertools.product(*[range(0, s, r) for s, r in zip(shape, region)]):
+ yield tuple(slice(st, builtins.min(st + r, s))
+ for st, r, s in zip(start, region, shape))
+
+
+def histogram(array, bins=256, range=None, strip_bytes=_DEFAULT_STRIP_BYTES, device=None):
+ """Streaming histogram over the whole array -- the substrate for global thresholds.
+
+ A histogram is an associative, memory-bound reduction: it equals the sum of per-region
+ histograms sharing the same bins, and the output size is fixed (``bins``) regardless of
+ array size. Returns ``(counts, bin_edges)`` like ``numpy.histogram``, and matches it
+ exactly. ``bins`` may be an int (with an optional ``range`` (lo, hi); if omitted, the
+ data min/max are found in one streaming pass) or a precomputed edges array. For a
+ per-channel/plane histogram, slice first: ``histogram(da[channel])``. ``device``
+ (None=inherit, 'cpu', 'cuda') runs the accumulation on that device; the result is host.
+ """
+ from ..utils import parse_dtype
+ dev = resolve_device(device)
+ itemsize = parse_dtype(array.dtype)[0].itemsize
+ bins_is_edges = not np.isscalar(bins)
+ if not bins_is_edges and range is None:
+ lo = float(min(array, device=device).compute()) # streaming min/max (memory-bound)
+ hi = float(max(array, device=device).compute())
+ if not (np.isfinite(lo) and np.isfinite(hi)) or lo == hi:
+ hi = lo + 1.0
+ range = (lo, hi)
+ if range is not None:
+ range = (float(range[0]), float(range[1]))
+
+ counts = None
+ edges = None
+ for region in _iter_region_slices(array.shape, itemsize, strip_bytes):
+ block = to_device(array._read_direct(region), dev)
+ xp = array_namespace(block)
+ if bins_is_edges:
+ c, edges = xp.histogram(block, bins=xp.asarray(bins))
+ else:
+ c, edges = xp.histogram(block, bins=bins, range=range)
+ counts = c if counts is None else counts + c
+ if counts is None: # empty array
+ edges = np.asarray(bins, dtype=float) if bins_is_edges else \
+ np.linspace(range[0], range[1], int(bins) + 1)
+ counts = np.zeros(len(edges) - 1, dtype=np.int64)
+ # bring back to host (counts/edges may be cupy) and return like numpy.histogram
+ return asnumpy(counts).astype(np.int64), asnumpy(edges)
+
+
+def unique(array, strip_bytes=_DEFAULT_STRIP_BYTES, device=None):
+ """Streaming distinct values over the whole array (like ``numpy.unique``: a sorted 1-D
+ array of the distinct values). Memory-bounded by the running set of distinct values plus
+ one region, so it is cheap when there are few distinct values (e.g. a label image) and
+ grows with that count otherwise. Eager, like ``histogram``: returns a host numpy array.
+ """
+ from ..utils import parse_dtype
+ dev = resolve_device(device)
+ dt = parse_dtype(array.dtype)[0]
+ if 0 in tuple(array.shape): # empty array -> no distinct values
+ return np.array([], dtype=dt)
+ acc = None
+ for region in _iter_region_slices(array.shape, dt.itemsize, strip_bytes):
+ block = to_device(array._read_direct(region), dev)
+ xp = array_namespace(block)
+ u = xp.unique(block)
+ acc = u if acc is None else xp.unique(xp.concatenate([acc, u]))
+ if acc is None: # empty array
+ return np.array([], dtype=dt)
+ return asnumpy(acc)
+
+
+__all__ = [
+ "ReduceTransform", "reduce",
+ "min", "max", "sum", "prod", "mean", "any", "all",
+ "var", "std", "argmin", "argmax", "median", "histogram", "unique",
+]
diff --git a/src/dyna_zarr/operations/scan.py b/src/dyna_zarr/operations/scan.py
new file mode 100644
index 0000000..07ac529
--- /dev/null
+++ b/src/dyna_zarr/operations/scan.py
@@ -0,0 +1,205 @@
+"""Prefix-scan ops along a single axis: cumsum / cumprod / cummax / cummin.
+
+A scan is neither pointwise nor a reduction: the output at index ``p`` along the scan axis
+depends on ALL input up to ``p``. Two execution paths, with different memory behavior:
+
+- ``io.write`` routes to ``scan_write`` (below), which streams with a BOUNDED CARRY: tile the
+ cross-section, walk the scan axis in strips keeping a running accumulator. Peak memory is
+ ~ one strip x one cross-section tile, independent of the scan axis and the array size --
+ FULLY memory-bounded, read-once/write-once. This is the path that matters.
+- The lazy ``read(key)`` (for ``compute`` / sub-slices) can't carry state between independent
+ reads, so it reads the input PREFIX ``[0, stop)`` along the scan axis, accumulates, and
+ slices back. That is correct and chunk-invariant, but bounded only by the prefix (~ the
+ scan axis), so prefer ``io.write`` for large scans.
+"""
+
+from __future__ import annotations
+
+from itertools import product as _iproduct
+
+import numpy as np
+
+from ._base import Transform
+from ._backend import array_namespace, asnumpy, resolve_device, to_device
+from .structural import _norm_key
+
+_DEFAULT_SCAN_MEM = 256 * 1024 * 1024 # per-call streaming budget (bytes)
+
+
+def _cumsum(xp, b, ax):
+ return xp.cumsum(b, axis=ax)
+
+
+def _cumprod(xp, b, ax):
+ return xp.cumprod(b, axis=ax)
+
+
+def _accumulate_via(ufunc_name):
+ """cummax/cummin via ``ufunc.accumulate``; falls back to host numpy if the array
+ namespace (e.g. some cupy versions) does not implement ufunc.accumulate."""
+ def acc(xp, b, ax):
+ uf = getattr(xp, ufunc_name)
+ try:
+ return uf.accumulate(b, axis=ax)
+ except (AttributeError, TypeError, NotImplementedError):
+ host = np.asarray(asnumpy(b))
+ return xp.asarray(getattr(np, ufunc_name).accumulate(host, axis=ax))
+ return acc
+
+
+_ACC = {
+ "cumsum": _cumsum,
+ "cumprod": _cumprod,
+ "cummax": _accumulate_via("maximum"),
+ "cummin": _accumulate_via("minimum"),
+}
+
+
+class ScanTransform(Transform):
+ """Lazy prefix scan along one axis. See the module docstring."""
+
+ def __init__(self, array, op, axis, device=None, name=None):
+ super().__init__()
+ self.array = array
+ self.op = op
+ self.acc = _ACC[op]
+ self.axis = int(axis) if axis >= 0 else array.ndim + int(axis)
+ if not (0 <= self.axis < array.ndim):
+ raise ValueError(f"scan axis {axis} out of range for ndim {array.ndim}")
+ self.device = device
+ self.name = name or op
+ self.shape = tuple(array.shape)
+ # dtype as numpy's accumulate would produce for this input dtype (e.g. cumsum keeps
+ # the integer dtype, matching numpy/dask -- users can astype to widen if needed).
+ self.dtype = np.asarray(self.acc(np, np.zeros((1,), dtype=array.dtype), 0)).dtype
+ self.chunks = None
+
+ def read(self, key):
+ norm = _norm_key(key, self.shape) # per axis (is_int, start, stop, step)
+ a = self.axis
+ ndim = len(self.shape)
+
+ # Read the full prefix on the scan axis; keep every axis (size-1 for int indices) so
+ # the scan axis stays at position ``a`` with no dim-shift bookkeeping.
+ read_key = []
+ for d in range(ndim):
+ is_int, start, stop, step = norm[d]
+ if d == a:
+ read_key.append(slice(0, stop)) # prefix [0, stop), step 1
+ elif is_int:
+ read_key.append(slice(start, start + 1)) # keepdim; squeezed at the end
+ else:
+ read_key.append(slice(start, stop, step))
+
+ dev = resolve_device(self.device)
+ block = to_device(self.array._read_direct(tuple(read_key)), dev)
+ xp = array_namespace(block)
+ acc = self.acc(xp, block, a)
+
+ # Select the requested output range on the scan axis (within the prefix), then drop
+ # every axis that was an integer index in the original key.
+ is_int_a, start_a, stop_a, step_a = norm[a]
+ sel = [slice(None)] * ndim
+ sel[a] = slice(start_a, start_a + 1) if is_int_a else slice(start_a, stop_a, step_a)
+ acc = acc[tuple(sel)]
+ drop = tuple(d for d in range(ndim) if norm[d][0])
+ return xp.squeeze(acc, axis=drop) if drop else acc
+
+
+# op -> (local scan, cross-strip combine with the running carry). carry is size-1 on the scan
+# axis and broadcasts over the strip; this makes the scan a bounded-carry stream.
+_COMBINE = {
+ "cumsum": lambda xp, carry, local: carry + local,
+ "cumprod": lambda xp, carry, local: carry * local,
+ "cummax": lambda xp, carry, local: xp.maximum(carry, local),
+ "cummin": lambda xp, carry, local: xp.minimum(carry, local),
+}
+
+
+def scan_write(source, op, axis, output_path, output_chunks=None,
+ max_mem=_DEFAULT_SCAN_MEM, dtype=None, zarr_format=2):
+ """Stream a prefix scan of ``source`` along ``axis`` to ``output_path``, FULLY
+ memory-bounded. A scan has a bounded carry (a single slab perpendicular to the scan
+ axis), so we tile the cross-section and, within each tile, walk the scan axis in strips
+ keeping a running carry: ``out_strip = combine(carry, local_scan(strip))``, then
+ ``carry = last slab of out_strip``. Peak memory is ~ one strip x one cross-section tile,
+ independent of the scan axis and of the array size. Read-once / write-once.
+ """
+ import zarr
+ from ..utils import parse_dtype
+
+ shape = tuple(int(s) for s in source.shape)
+ ndim = len(shape)
+ a = int(axis) if axis >= 0 else ndim + int(axis)
+ accfun = _ACC[op]
+ combine = _COMBINE[op]
+ dt = parse_dtype(dtype if dtype is not None else source.dtype)[0]
+ budget = max(1, int(max_mem) // dt.itemsize)
+
+ oc = tuple(output_chunks) if output_chunks is not None else (
+ tuple(source.chunks) if source.chunks else tuple(min(s, 256) for s in shape))
+ out = zarr.open(str(output_path), mode="w", shape=shape, chunks=oc,
+ dtype=dt, zarr_format=zarr_format)
+
+ # cross-section tile (all axes but the scan axis) sized to <= budget/2 so the carry fits
+ # and a strip of >= 2 slabs also fits; trailing axes grow first for contiguity.
+ cross_budget = max(1, budget // 2)
+ tile = [1] * ndim
+ acc = 1
+ for d in reversed([d for d in range(ndim) if d != a]):
+ tile[d] = min(shape[d], max(1, cross_budget // acc))
+ acc *= tile[d]
+ if tile[d] < shape[d]:
+ break
+ cross_size = 1
+ for d in range(ndim):
+ if d != a:
+ cross_size *= tile[d]
+ strip_a = max(1, budget // max(1, cross_size)) # slabs of the scan axis per read
+
+ cross_ranges = [range(0, shape[d], tile[d]) for d in range(ndim) if d != a]
+ other_axes = [d for d in range(ndim) if d != a]
+ for cross_origin in _iproduct(*cross_ranges):
+ base = [None] * ndim
+ for d, o in zip(other_axes, cross_origin):
+ base[d] = slice(o, min(o + tile[d], shape[d]))
+ carry = None
+ for a0 in range(0, shape[a], strip_a):
+ sl = list(base)
+ sl[a] = slice(a0, min(a0 + strip_a, shape[a]))
+ sl = tuple(sl)
+ block = np.asarray(asnumpy(source._read_direct(sl)))
+ local = accfun(np, block, a)
+ out_block = local if carry is None else combine(np, carry, local)
+ out[sl] = out_block
+ last = [slice(None)] * ndim
+ last[a] = slice(-1, None) # keepdims last slab -> next carry
+ carry = out_block[tuple(last)]
+ return output_path
+
+
+def _scan(array, op, axis, device=None):
+ return array._with_transform(ScanTransform(array, op, axis, device=device))
+
+
+def cumsum(array, axis, device=None):
+ """Cumulative sum along ``axis`` (shape-preserving)."""
+ return _scan(array, "cumsum", axis, device=device)
+
+
+def cumprod(array, axis, device=None):
+ """Cumulative product along ``axis`` (shape-preserving)."""
+ return _scan(array, "cumprod", axis, device=device)
+
+
+def cummax(array, axis, device=None):
+ """Cumulative maximum along ``axis`` (shape-preserving)."""
+ return _scan(array, "cummax", axis, device=device)
+
+
+def cummin(array, axis, device=None):
+ """Cumulative minimum along ``axis`` (shape-preserving)."""
+ return _scan(array, "cummin", axis, device=device)
+
+
+__all__ = ["ScanTransform", "cumsum", "cumprod", "cummax", "cummin"]
diff --git a/src/dyna_zarr/operations.py b/src/dyna_zarr/operations/structural.py
similarity index 53%
rename from src/dyna_zarr/operations.py
rename to src/dyna_zarr/operations/structural.py
index deb6436..bfee382 100644
--- a/src/dyna_zarr/operations.py
+++ b/src/dyna_zarr/operations/structural.py
@@ -1,11 +1,5 @@
-"""
-Array operations and transformations for DynamicArray.
-
-This module provides:
-- Transform classes for lazy operations
-- operations class with static methods for array manipulations
-- Support for all major numpy/dask.array operations
-"""
+"""Structural / coordinate transforms: concat, stack, slice, transpose, reshape,
+squeeze, flatten, pad, tile, roll, flip, expand_dims, swap_axes (+ slice_array)."""
import numpy as np
from typing import Tuple, Union, List, Optional, Any, TYPE_CHECKING
@@ -14,18 +8,70 @@
from dyna_zarr.dynamic_array import DynamicArray
-class Transform:
- """
- Base class for lazy transformations.
- """
+from ._base import Transform, _is_int_index, _perm_on_surviving
+from ._backend import array_namespace
- def __init__(self):
- self.shape = None
- self.chunks = None
- self.dtype = None
- def read(self, key):
- raise NotImplementedError
+def _reshape_read(array, new_shape, key):
+ """Windowed read for reshape/flatten (a C-order flat re-index).
+
+ An output region's elements occupy flat indices [fmin, fmax] (same in input and output
+ C-order). We read only the covering rectangle of input rows spanning that flat range,
+ then gather the region's flat positions from it. Memory ~= region + one inner slab, so
+ contiguous reads (io.write / compute) are memory-bound; only scatter patterns (a stepped
+ output) can widen the flat span toward the whole array.
+ """
+ axinfo = _norm_key(key, new_shape)
+ nnd = len(new_shape)
+ nstride = [1] * nnd
+ for a in range(nnd - 2, -1, -1):
+ nstride[a] = nstride[a + 1] * new_shape[a + 1]
+
+ kept_axes = [a for a, (is_int, *_r) in enumerate(axinfo) if not is_int]
+ kept_shape = tuple(len(range(axinfo[a][1], axinfo[a][2], axinfo[a][3])) for a in kept_axes)
+ flat = np.zeros(kept_shape, dtype=np.int64) # nd flat-index of each output elem
+ for a, (is_int, start, stop, step) in enumerate(axinfo):
+ st = int(nstride[a])
+ if is_int:
+ flat = flat + start * st
+ else:
+ coord = np.arange(start, stop, step, dtype=np.int64) * st
+ shp = [1] * len(kept_shape)
+ shp[kept_axes.index(a)] = coord.shape[0]
+ flat = flat + coord.reshape(shp)
+
+ in_shape = array.shape
+ in_stride0 = 1
+ for s in in_shape[1:]:
+ in_stride0 *= s
+ if flat.size == 0:
+ block = array._read_direct((slice(0, 0),) + (slice(None),) * (len(in_shape) - 1))
+ return array_namespace(block).reshape(block, kept_shape)
+ fmin, fmax = int(flat.min()), int(flat.max())
+ c0, c1 = fmin // in_stride0, fmax // in_stride0
+ box = array._read_direct((slice(c0, c1 + 1),) + (slice(None),) * (len(in_shape) - 1))
+ xp = array_namespace(box)
+ flatbox = box.reshape(-1)
+ idx = xp.asarray(flat - c0 * in_stride0) # into flatbox, on the box's device
+ return flatbox[idx]
+
+
+def _norm_key(key, shape):
+ """Per output axis -> (is_int, start, stop, step) with concrete non-negative values.
+ Integer indices become (True, idx, idx+1, 1); the caller squeezes those axes."""
+ if not isinstance(key, tuple):
+ key = (key,)
+ key = key + (slice(None),) * (len(shape) - len(key))
+ out = []
+ for a, k in enumerate(key):
+ size = shape[a]
+ if _is_int_index(k):
+ idx = int(k) if k >= 0 else size + int(k)
+ out.append((True, idx, idx + 1, 1))
+ else:
+ start, stop, step = k.indices(size)
+ out.append((False, start, stop, step))
+ return out
class ConcatenateTransform(Transform):
@@ -87,12 +133,12 @@ def read(self, key):
raise NotImplementedError(f"Indexing with {type(k)} not supported")
axis_slice = normalized_key[self.axis]
- start = axis_slice.start if axis_slice.start is not None else 0
- stop = axis_slice.stop if axis_slice.stop is not None else self.shape[self.axis]
- step = axis_slice.step if axis_slice.step is not None else 1
+ start, stop, step = axis_slice.indices(self.shape[self.axis])
- if step != 1:
- raise NotImplementedError("Step slicing not supported yet")
+ if step < 0:
+ raise NotImplementedError("Negative-step slicing on the concat axis not supported yet")
+ # Route the contiguous span [start, stop); the step is re-applied to the
+ # assembled result below so we read each source array only once.
# Find which arrays we need to read from
arrays_to_read = []
@@ -116,7 +162,12 @@ def read(self, key):
# OPTIMIZATION: Handle based on number of arrays
if len(arrays_to_read) == 0:
- raise ValueError("No data to read")
+ # Empty selection on the concat axis: return a correctly-shaped empty array
+ # (reading the first source with a zero-length concat-axis slice fixes the
+ # other axis sizes and the dtype without materialising any data).
+ empty_key = list(normalized_key)
+ empty_key[self.axis] = slice(0, 0)
+ result = self.arrays[0]._read_direct(tuple(empty_key))
elif len(arrays_to_read) == 1:
# Single array - no concatenation needed
result = arrays_to_read[0][0]._read_direct(arrays_to_read[0][1])
@@ -131,8 +182,13 @@ def read(self, key):
result = np.concatenate(result_parts, axis=self.axis)
+ # Re-apply the step on the concat axis (routing above read the contiguous span).
+ if step != 1:
+ step_idx = (slice(None),) * self.axis + (slice(None, None, step),)
+ result = result[step_idx]
+
# Remove dimensions that were indexed with int
- squeeze_axes = [i for i, k in enumerate(key) if isinstance(k, int)]
+ squeeze_axes = [i for i, k in enumerate(key) if _is_int_index(k)]
for ax in reversed(squeeze_axes):
result = np.squeeze(result, axis=ax)
@@ -173,16 +229,23 @@ def read(self, key):
source_key = key[:self.axis] + key[self.axis + 1:]
# Determine which arrays to read
- if isinstance(new_axis_key, int):
- # Single array
+ if _is_int_index(new_axis_key):
+ # Single array - the new axis is dropped
return self.arrays[new_axis_key]._read_direct(source_key)
- elif isinstance(new_axis_key, slice):
- start = new_axis_key.start if new_axis_key.start is not None else 0
- stop = new_axis_key.stop if new_axis_key.stop is not None else len(self.arrays)
- step = new_axis_key.step if new_axis_key.step is not None else 1
- parts = [self.arrays[i]._read_direct(source_key) for i in range(start, stop, step)]
- return np.stack(parts, axis=self.axis)
+ # Slice on the new axis. Insert the stacked axis at its position among the
+ # *surviving* source axes (earlier integer indices dropped their axes).
+ idxs = range(*new_axis_key.indices(len(self.arrays)))
+ insert_pos = sum(1 for j in range(self.axis) if not _is_int_index(key[j]))
+
+ if len(idxs) == 0:
+ # Empty selection: build an empty array of the correct per-part shape.
+ sample = self.arrays[0]._read_direct(source_key)
+ empty_shape = sample.shape[:insert_pos] + (0,) + sample.shape[insert_pos:]
+ return np.empty(empty_shape, dtype=sample.dtype)
+
+ parts = [self.arrays[i]._read_direct(source_key) for i in idxs]
+ return np.stack(parts, axis=insert_pos)
class SliceTransform(Transform):
@@ -193,58 +256,51 @@ class SliceTransform(Transform):
def __init__(self, array: 'DynamicArray', key):
super().__init__()
self.array = array
-
+
# Normalize key to tuple
if not isinstance(key, tuple):
key = (key,)
-
- self.key = key
- # Compute output shape and track new axes
+ chunks = array.chunks if array.chunks is not None else (1,) * array.ndim
new_shape = []
new_chunks = []
- original_dim = 0
- chunks = array.chunks if array.chunks is not None else (1,) * array.ndim
-
- # Track which original dimensions are being kept
- kept_dims = []
+ normalized = [] # store CONCRETE, non-negative keys so shape/read math is
+ original_dim = 0 # simple (negatives + None resolved once, here)
for k in key:
if k is np.newaxis:
- # Add new axis with chunk size 1
+ normalized.append(k)
new_shape.append(1)
new_chunks.append(1)
+ continue
+ if original_dim >= array.ndim:
+ raise IndexError("Too many indices for array")
+ size = array.shape[original_dim]
+ if isinstance(k, (int, np.integer)):
+ ki = int(k) if k >= 0 else size + int(k) # normalize negative index
+ if not (0 <= ki < size):
+ raise IndexError(
+ f"index {k} out of bounds for axis {original_dim} of size {size}")
+ normalized.append(ki) # int removes this dimension
+ elif isinstance(k, slice):
+ s = slice(*k.indices(size)) # resolves None + negatives
+ normalized.append(s)
+ new_shape.append(len(range(s.start, s.stop, s.step)))
+ new_chunks.append(chunks[original_dim])
else:
- if original_dim >= array.ndim:
- raise IndexError("Too many indices for array")
-
- if isinstance(k, int):
- # This dimension will be removed
- pass
- elif isinstance(k, slice):
- # Calculate size for this dimension
- size = array.shape[original_dim]
- start = k.start if k.start is not None else 0
- stop = k.stop if k.stop is not None else size
- step = k.step if k.step is not None else 1
- new_shape.append((stop - start + step - 1) // step)
- new_chunks.append(chunks[original_dim])
- kept_dims.append(original_dim)
- else:
- raise TypeError(f"Invalid index type: {type(k)}")
-
- original_dim += 1
+ raise TypeError(f"Invalid index type: {type(k)}")
+ original_dim += 1
- # Add remaining dimensions
+ # Add remaining (untouched) dimensions
for i in range(original_dim, array.ndim):
new_shape.append(array.shape[i])
new_chunks.append(chunks[i])
- kept_dims.append(i)
+ self.key = tuple(normalized)
self.shape = tuple(new_shape)
self.chunks = tuple(new_chunks)
self.dtype = array.dtype
- self.new_axes = [i for i, k in enumerate(key) if k is np.newaxis]
+ self.new_axes = [i for i, k in enumerate(self.key) if k is np.newaxis]
def read(self, read_key):
"""
@@ -279,11 +335,19 @@ def read(self, read_key):
# Dimension was sliced in the stored slice
read_elem = read_key[output_dim]
- # Get parameters of the stored slice
+ # Normalize the incoming read element (resolve None + negatives) against
+ # this output axis, so the composition below sees concrete non-negatives.
+ osize = self.shape[output_dim]
+ if isinstance(read_elem, (int, np.integer)):
+ read_elem = int(read_elem) if read_elem >= 0 else osize + int(read_elem)
+ elif isinstance(read_elem, slice):
+ read_elem = slice(*read_elem.indices(osize))
+
+ # Get parameters of the (already concrete) stored slice
orig_size = self.array.shape[input_dim]
- stored_start = stored_key_elem.start if stored_key_elem.start is not None else 0
- stored_stop = stored_key_elem.stop if stored_key_elem.stop is not None else orig_size
- stored_step = stored_key_elem.step if stored_key_elem.step is not None else 1
+ stored_start = stored_key_elem.start
+ stored_stop = stored_key_elem.stop
+ stored_step = stored_key_elem.step
if isinstance(read_elem, (int, np.integer)):
# Compose integer index with slice
@@ -319,10 +383,17 @@ def read(self, read_key):
else:
raise TypeError(f"Invalid stored key type: {type(stored_key_elem)}")
- # Add any remaining dimensions that weren't in the stored key
+ # Add any remaining dimensions that weren't in the stored key (normalize the
+ # read element against the underlying axis so no negatives reach the base array)
while input_dim < self.array.ndim:
if output_dim < len(read_key):
- full_key.append(read_key[output_dim])
+ elem = read_key[output_dim]
+ isize = self.array.shape[input_dim]
+ if isinstance(elem, (int, np.integer)):
+ elem = int(elem) if elem >= 0 else isize + int(elem)
+ elif isinstance(elem, slice):
+ elem = slice(*elem.indices(isize))
+ full_key.append(elem)
output_dim += 1
else:
full_key.append(slice(None))
@@ -371,18 +442,24 @@ def read(self, key):
# Pad key with full slices if needed
key = key + (slice(None),) * (len(self.shape) - len(key))
-
- # Build the key for the underlying array by removing the element at self.axis
- # since the underlying array doesn't have this dimension yet
+
+ # Element that applies to the inserted (singleton) axis, and the key for the
+ # underlying array (which lacks that axis).
+ key_ins = key[self.axis]
underlying_key = key[:self.axis] + key[self.axis + 1:]
-
+
# Read from underlying array
result = self.array._read_direct(underlying_key)
-
- # Add back the singleton dimension at the correct axis
- result = np.expand_dims(result, axis=self.axis)
-
- return result
+
+ # Insert the singleton axis at its position among the *surviving* axes: earlier
+ # axes indexed by an integer were dropped by the read, shifting the position.
+ insert_pos = sum(1 for j in range(self.axis) if not _is_int_index(key[j]))
+ result = np.expand_dims(result, axis=insert_pos)
+
+ # Apply the key element to the inserted axis (int drops it, slice sizes it 0/1).
+ idx = (slice(None),) * insert_pos + (key_ins,) + \
+ (slice(None),) * (result.ndim - insert_pos - 1)
+ return result[idx]
class SwapAxesTransform(Transform):
@@ -424,11 +501,12 @@ def read(self, key):
# Read from underlying array with unswapped key
result = self.array._read_direct(tuple(original_key))
-
- # Swap the axes in the result back
- result = np.swapaxes(result, self.axis1, self.axis2)
-
- return result
+
+ # A swap is a permutation; apply it on the surviving axes so integer indices
+ # that dropped an axis don't leave np.swapaxes with a stale axis index.
+ axes = list(range(self.array.ndim))
+ axes[self.axis1], axes[self.axis2] = axes[self.axis2], axes[self.axis1]
+ return _perm_on_surviving(result, tuple(key), tuple(axes))
class TransposeTransform(Transform):
@@ -475,10 +553,9 @@ def read(self, key):
# Read from underlying array with reordered key
result = self.array._read_direct(reordered_key)
- # Now transpose the result to match the expected output order
- result = np.transpose(result, self.axes)
-
- return result
+ # Transpose to output order, honouring any integer indices that dropped an
+ # axis (a plain np.transpose(result, self.axes) would use stale axis indices).
+ return _perm_on_surviving(result, key, self.axes)
# Extended operations from extended_operations.py
@@ -500,9 +577,7 @@ def __init__(self, array: 'DynamicArray', shape: Tuple[int, ...]):
self.dtype = array.dtype
def read(self, key):
- # For reshape, we need to materialize and reshape
- data = self.array._read_direct(slice(None))
- return data.reshape(self.new_shape)[key]
+ return _reshape_read(self.array, self.new_shape, key)
class SqueezeTransform(Transform):
@@ -562,17 +637,15 @@ def read(self, key):
if not isinstance(key, tuple):
key = (key,)
- # Build unsqueezed key by inserting slice(None) for squeezed axes
- # Track which positions in unsqueezed_key correspond to squeezed axes
+ # Build unsqueezed key by inserting slice(None) for squeezed axes.
unsqueezed_key = []
- axes_to_squeeze_in_result = []
+ squeezed_input_axes = []
output_idx = 0
for input_idx in range(self.array.ndim):
if input_idx in self.squeeze_axes:
- # This axis was squeezed - insert full slice
+ # This axis was squeezed - insert full slice (it is size 1 upstream).
unsqueezed_key.append(slice(None))
- # Mark this position for squeezing in the result
- axes_to_squeeze_in_result.append(len(unsqueezed_key) - 1)
+ squeezed_input_axes.append(input_idx)
else:
# This axis is preserved - use key from read operation
if output_idx < len(key):
@@ -580,17 +653,20 @@ def read(self, key):
else:
unsqueezed_key.append(slice(None))
output_idx += 1
-
+
# Read only the required region from underlying array
result = self.array._read_direct(tuple(unsqueezed_key))
-
-
- # Squeeze the marked axes (in reverse order to avoid index shifting)
- for ax in sorted(axes_to_squeeze_in_result, reverse=True):
- if ax < len(result.shape) and result.shape[ax] == 1:
- result = np.squeeze(result, axis=ax)
-
-
+
+ # A squeezed axis lands in the result at the position given by the number of
+ # *surviving* (non-integer-indexed) axes before it - integer indices earlier in
+ # the key drop their axes and shift everything left.
+ to_remove = [
+ sum(1 for j in range(input_idx) if not _is_int_index(unsqueezed_key[j]))
+ for input_idx in squeezed_input_axes
+ ]
+ for ax in sorted(to_remove, reverse=True):
+ result = np.squeeze(result, axis=ax)
+
return result
@@ -606,7 +682,7 @@ def __init__(self, array: 'DynamicArray'):
self.dtype = array.dtype
def read(self, key):
- return self.array._read_direct(slice(None)).flatten()[key]
+ return _reshape_read(self.array, self.shape, key)
class PadTransform(Transform):
@@ -632,8 +708,28 @@ def __init__(self, array: 'DynamicArray', pad_width: Union[int, Tuple]):
self.dtype = array.dtype
def read(self, key):
- data = self.array._read_direct(slice(None))
- return np.pad(data, self.pad_width)[key]
+ # Memory-bound: read only the core-overlapping input region, then pad just the
+ # border that this region includes (constant 0, like np.pad's default).
+ input_slices, pad_widths, crop, squeeze_axes = [], [], [], []
+ for a, (is_int, start, stop, step) in enumerate(_norm_key(key, self.shape)):
+ before, _after = self.pad_width[a]
+ in_size = self.array.shape[a]
+ core_lo, core_hi = max(start, before), min(stop, before + in_size)
+ input_slices.append(slice(core_lo - before, core_hi - before)
+ if core_hi > core_lo else slice(0, 0))
+ pad_before = max(0, min(stop, before) - start)
+ pad_after = max(0, stop - max(start, before + in_size))
+ pad_widths.append((pad_before, pad_after))
+ crop.append(slice(0, stop - start, step))
+ if is_int:
+ squeeze_axes.append(a)
+ core = self.array._read_direct(tuple(input_slices))
+ xp = array_namespace(core)
+ out = xp.pad(core, pad_widths) # constant 0
+ out = out[tuple(crop)]
+ for a in sorted(squeeze_axes, reverse=True):
+ out = xp.squeeze(out, axis=a)
+ return out
class TileTransform(Transform):
@@ -653,8 +749,20 @@ def __init__(self, array: 'DynamicArray', reps: Union[int, Tuple]):
self.dtype = array.dtype
def read(self, key):
- data = self.array._read_direct(slice(None))
- return np.tile(data, self.reps)[key]
+ # Bounded by the INPUT (the tile), not the full tiled OUTPUT: read the tile once,
+ # then gather the region via modulo indexing (out[i] = in[i % N] per axis).
+ axinfo = _norm_key(key, self.shape)
+ data = self.array._read_direct(tuple(slice(None) for _ in range(self.array.ndim)))
+ xp = array_namespace(data)
+ idx, squeeze_axes = [], []
+ for a, (is_int, start, stop, step) in enumerate(axinfo):
+ idx.append(xp.arange(start, stop, step) % self.array.shape[a])
+ if is_int:
+ squeeze_axes.append(a)
+ result = data[xp.ix_(*idx)]
+ for a in sorted(squeeze_axes, reverse=True):
+ result = xp.squeeze(result, axis=a)
+ return result
class RollTransform(Transform):
@@ -670,8 +778,36 @@ def __init__(self, array: 'DynamicArray', shift: int, axis: Optional[int] = None
self.dtype = array.dtype
def read(self, key):
- data = self.array._read_direct(slice(None))
- return np.roll(data, self.shift, axis=self.axis)[key]
+ if self.axis is None:
+ # flatten-roll: genuinely global. Rare; keep the whole-input fallback.
+ data = self.array._read_direct(tuple(slice(None) for _ in range(self.array.ndim)))
+ return np.roll(data, self.shift)[key]
+ # Memory-bound: out[i] = in[(i-shift) % N] along the axis, so a contiguous output
+ # run maps to 1 or 2 wrapped input segments. Read those, other axes contiguous,
+ # then crop step + squeeze ints.
+ A = self.axis if self.axis >= 0 else self.array.ndim + self.axis
+ N = self.array.shape[A]
+ shift = (self.shift % N) if N else 0
+ axinfo = _norm_key(key, self.shape)
+ base = [slice(start, stop) for (_ii, start, stop, _st) in axinfo]
+ _iiA, startA, stopA, _stA = axinfo[A]
+ L = stopA - startA
+ s0 = (startA - shift) % N if N else 0
+ if N == 0 or s0 + L <= N:
+ base[A] = slice(s0, s0 + L)
+ block = self.array._read_direct(tuple(base))
+ else:
+ k1, k2 = list(base), list(base)
+ k1[A] = slice(s0, N)
+ k2[A] = slice(0, s0 + L - N)
+ b1 = self.array._read_direct(tuple(k1))
+ b2 = self.array._read_direct(tuple(k2))
+ block = array_namespace(b1).concatenate([b1, b2], axis=A)
+ xp = array_namespace(block)
+ block = block[tuple(slice(0, stop - start, step) for (_ii, start, stop, step) in axinfo)]
+ for a in sorted((i for i, (ii, *_r) in enumerate(axinfo) if ii), reverse=True):
+ block = xp.squeeze(block, axis=a)
+ return block
class FlipTransform(Transform):
@@ -686,348 +822,138 @@ def __init__(self, array: 'DynamicArray', axis: int):
self.dtype = array.dtype
def read(self, key):
- data = self.array._read_direct(slice(None))
- return np.flip(data, axis=self.axis)[key]
-
+ # Memory-bound: read only the mirrored input window (positive-step slice, since
+ # zarr/tensorstore reject negative steps), flip it in-memory, then crop step +
+ # squeeze ints. out[start:stop) along the axis mirrors input[N-stop:N-start).
+ A = self.axis if self.axis >= 0 else self.array.ndim + self.axis
+ axinfo = _norm_key(key, self.shape)
+ input_slices = []
+ for a, (_is_int, start, stop, step) in enumerate(axinfo):
+ if a == A:
+ N = self.array.shape[a]
+ input_slices.append(slice(N - stop, N - start))
+ else:
+ input_slices.append(slice(start, stop))
+ block = self.array._read_direct(tuple(input_slices))
+ xp = array_namespace(block)
+ block = xp.flip(block, axis=A)
+ block = block[tuple(slice(0, stop - start, step) for (_ii, start, stop, step) in axinfo)]
+ for a in sorted((i for i, (ii, *_r) in enumerate(axinfo) if ii), reverse=True):
+ block = xp.squeeze(block, axis=a)
+ return block
-class ClipTransform(Transform):
- """Clip array values to a range."""
-
- def __init__(self, array: 'DynamicArray', a_min: Optional[float], a_max: Optional[float]):
- super().__init__()
- self.array = array
- self.a_min = a_min
- self.a_max = a_max
- self.shape = array.shape
- self.chunks = array.chunks
- self.dtype = array.dtype
-
- def read(self, key):
- data = self.array._read_direct(key)
- return np.clip(data, self.a_min, self.a_max)
-class AbsTransform(Transform):
- """Absolute value."""
-
- def __init__(self, array: 'DynamicArray'):
- super().__init__()
- self.array = array
- self.shape = array.shape
- self.chunks = array.chunks
- self.dtype = array.dtype
-
- def read(self, key):
- return np.abs(self.array._read_direct(key))
+def slice_array(array: 'DynamicArray', key) -> 'DynamicArray':
+ """Create a lazy slice of an array."""
+ transform = SliceTransform(array, key)
+ return array._with_transform(transform)
-class SignTransform(Transform):
- """Sign of array elements."""
-
- def __init__(self, array: 'DynamicArray'):
- super().__init__()
- self.array = array
- self.shape = array.shape
- self.chunks = array.chunks
- self.dtype = array.dtype
-
- def read(self, key):
- return np.sign(self.array._read_direct(key))
+# --------------------------------------------------------------------------- #
+# Public ops -- structural / coordinate (this module owns both transforms + ops)
+# --------------------------------------------------------------------------- #
+def expand_dims(array, axis):
+ """Add a new axis of length 1."""
+ return array._with_transform(ExpandDimsTransform(array, axis))
-class RoundTransform(Transform):
- """Round array elements."""
-
- def __init__(self, array: 'DynamicArray', decimals: int = 0):
- super().__init__()
- self.array = array
- self.decimals = decimals
- self.shape = array.shape
- self.chunks = array.chunks
- self.dtype = array.dtype
-
- def read(self, key):
- return np.round(self.array._read_direct(key), decimals=self.decimals)
+def concatenate(arrays, axis=0):
+ """Concatenate arrays along an existing axis."""
+ if not arrays:
+ raise ValueError("Need at least one array to concatenate")
+ return arrays[0]._with_transform(ConcatenateTransform(tuple(arrays), axis))
-class SqrtTransform(Transform):
- """Square root."""
-
- def __init__(self, array: 'DynamicArray'):
- super().__init__()
- self.array = array
- self.shape = array.shape
- self.chunks = array.chunks
- self.dtype = np.float64
-
- def read(self, key):
- return np.sqrt(self.array._read_direct(key))
+def stack(arrays, axis=0):
+ """Stack arrays along a new axis."""
+ if not arrays:
+ raise ValueError("Need at least one array to stack")
+ return arrays[0]._with_transform(StackTransform(tuple(arrays), axis))
-class WhereTransform(Transform):
- """Conditional element selection."""
-
- def __init__(self, condition: 'DynamicArray', x: 'DynamicArray', y: 'DynamicArray'):
- super().__init__()
- if not (condition.shape == x.shape == y.shape):
- raise ValueError("All arrays must have the same shape")
-
- self.condition = condition
- self.x = x
- self.y = y
- self.shape = x.shape
- self.chunks = x.chunks
- self.dtype = x.dtype
-
- def read(self, key):
- cond = self.condition._read_direct(key)
- x_data = self.x._read_direct(key)
- y_data = self.y._read_direct(key)
- return np.where(cond, x_data, y_data)
+def swap_axes(array, axis1, axis2):
+ """Swap two axes."""
+ return array._with_transform(SwapAxesTransform(array, axis1, axis2))
-class MultiplyTransform(Transform):
- """Element-wise multiplication."""
-
- def __init__(self, array1: 'DynamicArray', array2: Union['DynamicArray', float]):
- super().__init__()
- self.array1 = array1
- self.array2 = array2
- self.shape = array1.shape
- self.chunks = array1.chunks
- self.dtype = array1.dtype
-
- def read(self, key):
- # Import here to avoid circular dependency
- from dyna_zarr.dynamic_array import DynamicArray
-
- data1 = self.array1._read_direct(key)
- if isinstance(self.array2, DynamicArray):
- data2 = self.array2._read_direct(key)
- else:
- data2 = self.array2
- return data1 * data2
+def transpose(array, axes):
+ """Permute array dimensions."""
+ return array._with_transform(TransposeTransform(array, axes))
-class AddTransform(Transform):
- """Element-wise addition."""
-
- def __init__(self, array1: 'DynamicArray', array2: Union['DynamicArray', float]):
- super().__init__()
- self.array1 = array1
- self.array2 = array2
- self.shape = array1.shape
- self.chunks = array1.chunks
- self.dtype = array1.dtype
-
- def read(self, key):
- # Import here to avoid circular dependency
- from dyna_zarr.dynamic_array import DynamicArray
-
- data1 = self.array1._read_direct(key)
- if isinstance(self.array2, DynamicArray):
- data2 = self.array2._read_direct(key)
- else:
- data2 = self.array2
- return data1 + data2
+def reshape(array, shape):
+ """Reshape array to a new shape (C-order).
-class MinTransform(Transform):
- """Lazy minimum reduction along specified axes."""
- def __init__(self, array: 'DynamicArray', axis: Optional[int] = None):
- super().__init__()
- self.array = array
- self.axis = axis
-
- if axis is None:
- self.shape = ()
- self.chunks = None
- else:
- normalized_axis = axis if axis >= 0 else array.ndim + axis
- if normalized_axis < 0 or normalized_axis >= array.ndim:
- raise ValueError(f"axis {axis} out of bounds for dimension {array.ndim}")
- self.shape = array.shape[:normalized_axis] + array.shape[normalized_axis + 1:]
- self.chunks = array.chunks[:normalized_axis] + array.chunks[normalized_axis + 1:] if array.chunks else None
-
- self.dtype = array.dtype
-
- def read(self, key):
- """Read entire array and compute minimum."""
- full_data = self.array._read_direct(tuple(slice(None) for _ in range(self.array.ndim)))
- return np.min(full_data, axis=self.axis)
+ Memory: ``io.write`` streams this via the reindex writer (per output chunk, one input
+ chunk at a time), so writing a reshape is hard-bounded to ~1 input chunk + 1 output
+ chunk regardless of array size. A lazy ``compute()`` / sub-slice still uses a covering
+ read (heavier, since a flat re-index conflicts with nd chunk layout).
+ """
+ return array._with_transform(ReshapeTransform(array, shape))
-class MaxTransform(Transform):
- """Lazy maximum reduction along specified axes."""
- def __init__(self, array: 'DynamicArray', axis: Optional[int] = None):
- super().__init__()
- self.array = array
- self.axis = axis
-
- if axis is None:
- self.shape = ()
- self.chunks = None
- else:
- normalized_axis = axis if axis >= 0 else array.ndim + axis
- if normalized_axis < 0 or normalized_axis >= array.ndim:
- raise ValueError(f"axis {axis} out of bounds for dimension {array.ndim}")
- self.shape = array.shape[:normalized_axis] + array.shape[normalized_axis + 1:]
- self.chunks = array.chunks[:normalized_axis] + array.chunks[normalized_axis + 1:] if array.chunks else None
-
- self.dtype = array.dtype
-
- def read(self, key):
- """Read entire array and compute maximum."""
- full_data = self.array._read_direct(tuple(slice(None) for _ in range(self.array.ndim)))
- return np.max(full_data, axis=self.axis)
+def squeeze(array, axis=None):
+ """Remove singleton dimensions."""
+ return array._with_transform(SqueezeTransform(array, axis))
-# operations class with static methods for creating transforms
+def flatten(array):
+ """Flatten array to 1D (C-order).
-class operations:
- """
- Array operations following numpy API conventions.
- All methods return lazy DynamicArray objects with transforms applied.
+ Memory: ``io.write`` streams this via the reindex writer, hard-bounded to ~1 input
+ chunk + 1 output chunk regardless of array size (verified flat: 170MB->679MB arrays
+ all peak ~135MB). A lazy ``compute()`` still materializes.
"""
-
- @staticmethod
- def expand_dims(array: 'DynamicArray', axis: int) -> 'DynamicArray':
- """Add a new axis of length 1."""
- transform = ExpandDimsTransform(array, axis)
- return array._with_transform(transform)
-
- @staticmethod
- def concatenate(arrays: List['DynamicArray'], axis: int = 0) -> 'DynamicArray':
- """Concatenate arrays along an existing axis."""
- if not arrays:
- raise ValueError("Need at least one array to concatenate")
- transform = ConcatenateTransform(tuple(arrays), axis)
- return arrays[0]._with_transform(transform)
-
- @staticmethod
- def stack(arrays: List['DynamicArray'], axis: int = 0) -> 'DynamicArray':
- """Stack arrays along a new axis."""
- if not arrays:
- raise ValueError("Need at least one array to stack")
- transform = StackTransform(tuple(arrays), axis)
- return arrays[0]._with_transform(transform)
-
- @staticmethod
- def swap_axes(array: 'DynamicArray', axis1: int, axis2: int) -> 'DynamicArray':
- """Swap two axes."""
- transform = SwapAxesTransform(array, axis1, axis2)
- return array._with_transform(transform)
-
- @staticmethod
- def transpose(array: 'DynamicArray', axes: Tuple[int, ...]) -> 'DynamicArray':
- """Permute array dimensions."""
- transform = TransposeTransform(array, axes)
- return array._with_transform(transform)
-
- @staticmethod
- def reshape(array: 'DynamicArray', shape: Tuple[int, ...]) -> 'DynamicArray':
- """Reshape array to new shape."""
- transform = ReshapeTransform(array, shape)
- return array._with_transform(transform)
-
- @staticmethod
- def squeeze(array: 'DynamicArray', axis: Optional[int] = None) -> 'DynamicArray':
- """Remove singleton dimensions."""
- transform = SqueezeTransform(array, axis)
- return array._with_transform(transform)
-
- @staticmethod
- def flatten(array: 'DynamicArray') -> 'DynamicArray':
- """Flatten array to 1D."""
- transform = FlattenTransform(array)
- return array._with_transform(transform)
-
- @staticmethod
- def pad(array: 'DynamicArray', pad_width: Union[int, Tuple]) -> 'DynamicArray':
- """Pad array."""
- transform = PadTransform(array, pad_width)
- return array._with_transform(transform)
-
- @staticmethod
- def tile(array: 'DynamicArray', reps: Union[int, Tuple]) -> 'DynamicArray':
- """Repeat array along dimensions."""
- transform = TileTransform(array, reps)
- return array._with_transform(transform)
-
- @staticmethod
- def roll(array: 'DynamicArray', shift: int, axis: Optional[int] = None) -> 'DynamicArray':
- """Roll array elements along an axis."""
- transform = RollTransform(array, shift, axis)
- return array._with_transform(transform)
-
- @staticmethod
- def flip(array: 'DynamicArray', axis: int) -> 'DynamicArray':
- """Flip array along an axis."""
- transform = FlipTransform(array, axis)
- return array._with_transform(transform)
-
- @staticmethod
- def clip(array: 'DynamicArray', a_min: Optional[float], a_max: Optional[float]) -> 'DynamicArray':
- """Clip array values to a range."""
- transform = ClipTransform(array, a_min, a_max)
- return array._with_transform(transform)
-
- @staticmethod
- def abs(array: 'DynamicArray') -> 'DynamicArray':
- """Absolute value."""
- transform = AbsTransform(array)
- return array._with_transform(transform)
-
- @staticmethod
- def sign(array: 'DynamicArray') -> 'DynamicArray':
- """Sign of array elements."""
- transform = SignTransform(array)
- return array._with_transform(transform)
-
- @staticmethod
- def round(array: 'DynamicArray', decimals: int = 0) -> 'DynamicArray':
- """Round array elements."""
- transform = RoundTransform(array, decimals)
- return array._with_transform(transform)
-
- @staticmethod
- def sqrt(array: 'DynamicArray') -> 'DynamicArray':
- """Square root."""
- transform = SqrtTransform(array)
- return array._with_transform(transform)
-
- @staticmethod
- def where(condition: 'DynamicArray', x: 'DynamicArray', y: 'DynamicArray') -> 'DynamicArray':
- """Conditional element selection."""
- transform = WhereTransform(condition, x, y)
- return x._with_transform(transform)
-
- @staticmethod
- def multiply(array1: 'DynamicArray', array2: Union['DynamicArray', float]) -> 'DynamicArray':
- """Element-wise multiplication."""
- transform = MultiplyTransform(array1, array2)
- return array1._with_transform(transform)
-
- @staticmethod
- def add(array1: 'DynamicArray', array2: Union['DynamicArray', float]) -> 'DynamicArray':
- """Element-wise addition."""
- transform = AddTransform(array1, array2)
- return array1._with_transform(transform)
-
- @staticmethod
- def min(array: 'DynamicArray', axis: Optional[int] = None) -> 'DynamicArray':
- """Compute minimum along axis (lazy reduction)."""
- transform = MinTransform(array, axis)
- return array._with_transform(transform)
-
- @staticmethod
- def max(array: 'DynamicArray', axis: Optional[int] = None) -> 'DynamicArray':
- """Compute maximum along axis (lazy reduction)."""
- transform = MaxTransform(array, axis)
- return array._with_transform(transform)
+ return array._with_transform(FlattenTransform(array))
-def slice_array(array: 'DynamicArray', key) -> 'DynamicArray':
- """Create a lazy slice of an array."""
- transform = SliceTransform(array, key)
- return array._with_transform(transform)
+def pad(array, pad_width):
+ """Pad array."""
+ return array._with_transform(PadTransform(array, pad_width))
+
+
+def tile(array, reps):
+ """Repeat array along dimensions."""
+ return array._with_transform(TileTransform(array, reps))
+
+
+def roll(array, shift, axis=None):
+ """Roll array elements along an axis."""
+ return array._with_transform(RollTransform(array, shift, axis))
+
+
+def flip(array, axis):
+ """Flip array along an axis."""
+ return array._with_transform(FlipTransform(array, axis))
+
+
+def rot90(array, k=1, axes=(0, 1)):
+ """Rotate array by 90 degrees ``k`` times in the plane of ``axes`` (like numpy.rot90).
+ Composed from the validated flip/transpose transforms, so it stays lazy + backend-agnostic."""
+ ndim = array.ndim
+ a0 = axes[0] if axes[0] >= 0 else ndim + axes[0]
+ a1 = axes[1] if axes[1] >= 0 else ndim + axes[1]
+ if a0 == a1 or not (0 <= a0 < ndim and 0 <= a1 < ndim):
+ raise ValueError(f"invalid rotation axes {axes} for ndim {ndim}")
+ k %= 4
+ if k == 0:
+ return array
+ if k == 2:
+ return flip(flip(array, a0), a1)
+ perm = list(range(ndim))
+ perm[a0], perm[a1] = perm[a1], perm[a0]
+ if k == 1:
+ return transpose(flip(array, a1), tuple(perm))
+ return flip(transpose(array, tuple(perm)), a1) # k == 3
+
+
+__all__ = [
+ "ConcatenateTransform", "StackTransform", "SliceTransform", "ExpandDimsTransform",
+ "SwapAxesTransform", "TransposeTransform", "ReshapeTransform", "SqueezeTransform",
+ "FlattenTransform", "PadTransform", "TileTransform", "RollTransform", "FlipTransform",
+ "slice_array",
+ "expand_dims", "concatenate", "stack", "swap_axes", "transpose", "reshape",
+ "squeeze", "flatten", "pad", "tile", "roll", "flip", "rot90",
+]
diff --git a/src/dyna_zarr/rechunk.py b/src/dyna_zarr/rechunk.py
new file mode 100644
index 0000000..2e723b1
--- /dev/null
+++ b/src/dyna_zarr/rechunk.py
@@ -0,0 +1,233 @@
+"""Two-phase disk-staged rechunk (Rechunker's algorithm, dask-free).
+
+Rechunking source chunks -> target chunks directly can blow memory when the chunk grids
+don't align. Rechunker's fix: stage through an on-disk intermediate whose chunks divide
+BOTH grids, so each hop is chunk-aligned -> each source chunk read once, each target chunk
+written once, memory bounded to one chunk.
+
+ int_chunks[d] = gcd(source_chunks[d], target_chunks[d]) # divides both grids
+
+ - source divides target (source finer) -> one pass, consolidate (iterate target chunks).
+ - target divides source (target finer) -> one pass, split (iterate source chunks).
+ - otherwise -> two passes via a disk intermediate (int_chunks):
+ phase 1 source -> intermediate (iterate source chunks; int|source -> aligned)
+ phase 2 intermediate -> target (iterate target chunks; int|target -> aligned)
+
+Same-shape only (it's a re-chunk, not a re-shape). reshape/flatten build on top of this by
+rechunking to a flat-contiguous layout, then relabeling the shape.
+"""
+
+import math
+import shutil
+from concurrent.futures import ThreadPoolExecutor
+from itertools import product
+
+import numpy as np
+import zarr
+
+from .utils import parse_dtype
+from .operations._backend import asnumpy
+
+DEFAULT_MAX_MEM = 256 * 1024 * 1024 # per-worker region-buffer budget, in bytes
+DEFAULT_MAX_WORKERS = 4 # concurrent read->write regions
+
+
+def _read(src, sl):
+ from .dynamic_array import DynamicArray
+ if isinstance(src, DynamicArray):
+ return asnumpy(src._read_direct(sl))
+ return np.asarray(src[sl])
+
+
+def _expand_region(base_chunk, shape, budget):
+ """Grow a copy region outward from one ``base_chunk`` (the coarser grid's chunk) by WHOLE
+ chunk multiples -- innermost axis first (C-order, keeps reads contiguous) -- while the
+ element count stays within ``budget``. Result is chunk-aligned on both grids (so read-once
+ / write-once holds) and bounds each worker's buffer to ~budget. For tiny chunks this
+ batches many into one big I/O; for a chunk already >= budget it stays at one chunk (floor).
+ """
+ ndim = len(shape)
+ region = [min(int(base_chunk[d]), int(shape[d])) for d in range(ndim)]
+ for ax in range(ndim - 1, -1, -1):
+ while region[ax] < shape[ax]:
+ nxt = min(region[ax] + int(base_chunk[ax]), int(shape[ax]))
+ prod = 1
+ for d in range(ndim):
+ prod *= nxt if d == ax else region[d]
+ if prod <= budget:
+ region[ax] = nxt
+ else:
+ break
+ return tuple(region)
+
+
+def _copy(src, dst, base_chunk, shape, budget, max_workers):
+ """Copy src->dst region by region. Each region is a multiple of ``base_chunk`` (the COARSER
+ grid) sized to ``budget`` -> chunk-aligned both sides (read-once / write-once), disjoint ->
+ parallel-safe (distinct chunk files). Up to ``max_workers`` regions in flight => peak
+ ~= max_workers * region_bytes."""
+ ndim = len(shape)
+ region = _expand_region(base_chunk, shape, budget)
+ origins = list(product(*[range(0, shape[d], region[d]) for d in range(ndim)]))
+
+ def task(origin):
+ sl = tuple(slice(o, min(o + region[d], shape[d]))
+ for d, o in zip(range(ndim), origin))
+ dst[sl] = _read(src, sl)
+
+ if max_workers <= 1 or len(origins) <= 1:
+ for o in origins:
+ task(o)
+ else:
+ # ThreadPoolExecutor runs <= max_workers tasks concurrently; reads happen INSIDE the
+ # task, so at most max_workers regions are resident at once. list() surfaces exceptions.
+ with ThreadPoolExecutor(max_workers=max_workers) as ex:
+ list(ex.map(task, origins))
+
+
+def rechunk(source, target_chunks, output_path, dtype=None, zarr_format=2,
+ intermediate_path=None, max_mem=DEFAULT_MAX_MEM, max_workers=DEFAULT_MAX_WORKERS):
+ """Rechunk ``source`` (a DynamicArray or zarr, same output shape) to ``target_chunks`` at
+ ``output_path``, disk-staged and memory-bounded. Returns output_path.
+
+ Memory model (per-worker, matching io.write's region_size_mb + max_workers): each worker
+ buffers up to ``max_mem`` bytes (one region = as many whole aligned chunks as fit, >= one
+ chunk); ``max_workers`` regions run concurrently, so peak ~= ``max_workers * max_mem``.
+ """
+ shape = tuple(int(s) for s in source.shape)
+ ndim = len(shape)
+ sc = tuple(int(c) for c in (source.chunks or shape))
+ tc = tuple(int(c) for c in target_chunks)
+ if len(tc) != ndim:
+ raise ValueError(f"target_chunks {tc} rank != source rank {ndim}")
+ dt = parse_dtype(dtype if dtype is not None else source.dtype)[0]
+ budget = max(1, int(max_mem) // dt.itemsize) # region budget in elements
+
+ out = zarr.open(str(output_path), mode="w", shape=shape, chunks=tc,
+ dtype=dt, zarr_format=zarr_format)
+
+ src_divides_tgt = all(tc[d] % sc[d] == 0 for d in range(ndim)) # source finer
+ tgt_divides_src = all(sc[d] % tc[d] == 0 for d in range(ndim)) # target finer
+
+ if src_divides_tgt:
+ _copy(source, out, tc, shape, budget, max_workers) # consolidate: base = target chunk
+ elif tgt_divides_src:
+ _copy(source, out, sc, shape, budget, max_workers) # split: base = source chunk
+ else:
+ ic = tuple(math.gcd(sc[d], tc[d]) for d in range(ndim))
+ ipath = intermediate_path or (str(output_path) + ".rechunk_int.zarr")
+ interm = zarr.open(ipath, mode="w", shape=shape, chunks=ic, dtype=dt,
+ zarr_format=zarr_format)
+ _copy(source, interm, sc, shape, budget, max_workers) # phase 1 (int | source)
+ _copy(interm, out, tc, shape, budget, max_workers) # phase 2 (int | target)
+ try:
+ shutil.rmtree(ipath)
+ except Exception:
+ pass
+ return output_path
+
+
+def flatten_write(source, output_path, output_chunks=None, max_mem=DEFAULT_MAX_MEM,
+ max_workers=DEFAULT_MAX_WORKERS, dtype=None, zarr_format=2):
+ """Flatten ``source`` (nd) to 1D at ``output_path``, memory-bounded + read-once.
+
+ Two steps: (1) rechunk source to a FLAT-CONTIGUOUS layout (read-once source via the
+ rechunk engine); (2) relabel those flat-contiguous slabs to the 1D output (read-once).
+ The nd-chunk -> flat-order reorder happens inside the disk-staged rechunk, so RAM stays
+ bounded to ~one flat-contiguous unit regardless of array size.
+
+ Sizing the flat-contiguous unit to ``max_mem``: a chunk is a single C-order contiguous
+ flat run iff every axis AFTER some split axis ``a`` is full and every axis BEFORE is 1,
+ i.e. ``(1,..,1, ca, D_{a+1}..D_last)`` -> ``ca * suffix[a]`` contiguous elements
+ (``suffix[a] = prod(shape[a+1:])``). Pick the shallowest ``a`` whose full-trailing block
+ ``suffix[a]`` already fits the budget, then take as many ``ca`` steps as fit. This bounds
+ the unit to ``max_mem`` even when a whole trailing row (split axis 0) would blow it.
+ """
+ shape = tuple(int(s) for s in source.shape)
+ ndim = len(shape)
+ dt = parse_dtype(dtype if dtype is not None else source.dtype)[0]
+ budget = max(1, int(max_mem) // dt.itemsize) # in elements
+
+ suffix = [1] * ndim # suffix[d] = prod(shape[d+1:])
+ for d in range(ndim - 2, -1, -1):
+ suffix[d] = suffix[d + 1] * shape[d + 1]
+
+ a = ndim - 1 # split axis (deepest = last)
+ for d in range(ndim):
+ if suffix[d] <= budget: # trailing block already fits
+ a = d
+ break
+ ca = max(1, min(shape[a], budget // suffix[a])) # steps of the split axis that fit
+ ff_chunks = tuple(1 if d < a else (ca if d == a else shape[d]) for d in range(ndim))
+
+ ff_path = str(output_path) + ".flatten_ff.zarr"
+ rechunk(source, ff_chunks, ff_path, dtype=dt, zarr_format=zarr_format,
+ max_mem=max_mem, max_workers=max_workers) # read-once source, parallel
+ ff = zarr.open(ff_path, mode="r")
+
+ N = suffix[0] * shape[0]
+ unit = ca * suffix[a] # elements per flat-contiguous slab
+ oc = tuple(output_chunks) if output_chunks is not None else (min(N, unit),)
+ out = zarr.open(str(output_path), mode="w", shape=(N,), chunks=oc, dtype=dt,
+ zarr_format=zarr_format)
+
+ # Relabel flat-contiguous units to the 1D output, in C-order: leading axes one index at a
+ # time, split axis in steps of ca, trailing axes whole -> each unit is a contiguous output
+ # range. Kept SEQUENTIAL: user output_chunks need not align to units, so parallel writes
+ # could collide on a shared output chunk. Memory here = one unit (<= max_mem); it's mostly
+ # I/O, and the expensive reorder already happened (parallel) in the rechunk above.
+ lead_ranges = [range(shape[d]) for d in range(a)] + [range(0, shape[a], ca)]
+ for origin in product(*lead_ranges):
+ lead, ka = origin[:a], origin[a]
+ da = min(ca, shape[a] - ka)
+ sl = tuple([slice(i, i + 1) for i in lead] + [slice(ka, ka + da)]
+ + [slice(None)] * (ndim - a - 1))
+ block = np.asarray(ff[sl]) # one ff slab (read-once)
+ f0 = sum(lead[d] * suffix[d] for d in range(a)) + ka * suffix[a]
+ out[f0:f0 + da * suffix[a]] = block.reshape(-1)
+ shutil.rmtree(ff_path, ignore_errors=True)
+ return output_path
+
+
+def reshape_write(source, target_shape, output_path, output_chunks=None,
+ max_mem=DEFAULT_MAX_MEM, max_workers=DEFAULT_MAX_WORKERS,
+ dtype=None, zarr_format=2):
+ """Reshape ``source`` -> ``target_shape`` (C-order) at ``output_path``, memory-bounded.
+
+ Two staged steps: (1) flatten source to a 1D CONTIGUOUS ``F`` (read-once source, via the
+ disk-staged rechunk engine); (2) reindex ``F`` -> ``target_shape`` streaming per output
+ chunk. Feeding reindex a 1D-contiguous input is the key: reindex was only slow on nd
+ sources because a flat range forced whole nd chunks -- over contiguous ``F`` each output
+ chunk maps to contiguous ``F`` ranges, so there's no chunk-amplification. Peak memory ~
+ max(one flatten unit [<= max_mem], one F chunk + one output chunk of the reindex).
+
+ Shape change happens only at the F -> target flat-identity relabel (output flat index ==
+ input flat index); everything else is same-shape staging.
+ """
+ from .dynamic_array import DynamicArray
+ from .reindex import reindex_write
+
+ target_shape = tuple(int(s) for s in target_shape)
+ src_shape = tuple(int(s) for s in source.shape)
+ if int(np.prod(src_shape)) != int(np.prod(target_shape)):
+ raise ValueError(f"reshape needs equal size: {src_shape} -> {target_shape}")
+ dt = parse_dtype(dtype if dtype is not None else source.dtype)[0]
+
+ N = int(np.prod(target_shape))
+ if output_chunks is None:
+ output_chunks = ((min(N, 1 << 20),) if len(target_shape) == 1
+ else tuple(min(s, 256) for s in target_shape))
+ else:
+ output_chunks = tuple(int(c) for c in output_chunks)
+
+ f_path = str(output_path) + ".reshape_1d.zarr"
+ f_chunk = (min(N, 1 << 22),) # 1D F chunks (~16MB f32), lean reindex input
+ flatten_write(source, f_path, output_chunks=f_chunk, max_mem=max_mem,
+ max_workers=max_workers, dtype=dt, zarr_format=zarr_format)
+ F = DynamicArray(zarr.open(f_path, mode="r"))
+ try:
+ reindex_write(F, target_shape, output_path, output_chunks,
+ dtype=dt, zarr_format=zarr_format)
+ finally:
+ shutil.rmtree(f_path, ignore_errors=True)
+ return output_path
diff --git a/src/dyna_zarr/reindex.py b/src/dyna_zarr/reindex.py
new file mode 100644
index 0000000..035c632
--- /dev/null
+++ b/src/dyna_zarr/reindex.py
@@ -0,0 +1,104 @@
+"""Streaming reindex writer for reshape / flatten (a C-order flat re-index).
+
+These conflict with nd chunk layout, so the region writer either materializes the whole
+input or thrashes chunks. This engine instead is **output-chunk driven with per-input-chunk
+streaming**:
+
+ for each OUTPUT chunk O:
+ buffer = empty(O)
+ for each INPUT chunk that feeds O (read ONE AT A TIME, then released):
+ scatter its overlapping elements into buffer
+ write O once
+
+So each output chunk is produced by a single task (no write race, at any parallelism), and
+memory is hard-bounded to **one input chunk + one output chunk** (+ the O-sized index
+arrays), regardless of how many input chunks feed O or how far apart they are. The cost is
+read amplification -- an input chunk is re-read once per output chunk that needs it -- which
+an optional bounded LRU cache (``cache_size`` input chunks) trades back for memory.
+
+reshape/flatten are the flat-identity case: output flat index == input flat index.
+"""
+
+from itertools import product
+
+import numpy as np
+import zarr
+
+from .utils import parse_dtype
+from .operations._backend import asnumpy
+
+
+def _c_strides(shape):
+ strides = [1] * len(shape)
+ for a in range(len(shape) - 2, -1, -1):
+ strides[a] = strides[a + 1] * shape[a + 1]
+ return strides
+
+
+def reindex_write(input_array, out_shape, output_path, output_chunks,
+ dtype=None, zarr_format=2, cache_size=0):
+ """Write ``input_array`` reshaped (C-order) to ``out_shape`` at ``output_path``,
+ streaming per output chunk. Memory ~ one input chunk + one output chunk + O-sized index
+ arrays. ``cache_size`` (input chunks) optionally caches recent input chunks to cut
+ re-reads."""
+ in_shape = tuple(int(s) for s in input_array.shape)
+ out_shape = tuple(int(s) for s in out_shape)
+ if int(np.prod(in_shape)) != int(np.prod(out_shape)):
+ raise ValueError(f"reindex needs equal size: {in_shape} vs {out_shape}")
+ in_chunks = tuple(int(c) for c in (input_array.chunks or in_shape))
+ output_chunks = tuple(int(c) for c in output_chunks)
+ dt = parse_dtype(dtype if dtype is not None else input_array.dtype)[0]
+
+ out = zarr.open(str(output_path), mode="w", shape=out_shape,
+ chunks=output_chunks, dtype=dt, zarr_format=zarr_format)
+
+ in_ndim, out_ndim = len(in_shape), len(out_shape)
+ out_strides = _c_strides(out_shape)
+ grid = tuple((in_shape[d] + in_chunks[d] - 1) // in_chunks[d] for d in range(in_ndim))
+
+ cache = {} if cache_size > 0 else None
+ order = []
+
+ def read_input_chunk(cid, cc):
+ if cache is not None and cid in cache:
+ return cache[cid]
+ sl = tuple(slice(cc[d] * in_chunks[d],
+ min(cc[d] * in_chunks[d] + in_chunks[d], in_shape[d]))
+ for d in range(in_ndim))
+ block = np.asarray(asnumpy(input_array._read_direct(sl)))
+ if cache is not None:
+ cache[cid] = block
+ order.append(cid)
+ while len(order) > cache_size:
+ cache.pop(order.pop(0), None)
+ return block
+
+ for origin in product(*[range(0, out_shape[d], output_chunks[d]) for d in range(out_ndim)]):
+ o_sl = tuple(slice(origin[d], min(origin[d] + output_chunks[d], out_shape[d]))
+ for d in range(out_ndim))
+ o_shape = tuple(s.stop - s.start for s in o_sl)
+
+ # flat index (C-order in out_shape) of every element of this output chunk
+ flat = np.zeros(o_shape, dtype=np.int64)
+ for d in range(out_ndim):
+ coord = (o_sl[d].start + np.arange(o_shape[d], dtype=np.int64)) * np.int64(out_strides[d])
+ shp = [1] * out_ndim
+ shp[d] = o_shape[d]
+ flat += coord.reshape(shp)
+ flat = flat.ravel()
+
+ # same flat index in the INPUT -> input coords + which input chunk each belongs to
+ in_coords = np.unravel_index(flat, in_shape)
+ chunk_lin = np.ravel_multi_index(
+ tuple(in_coords[d] // in_chunks[d] for d in range(in_ndim)), grid)
+
+ obuf = np.empty(flat.shape[0], dtype=dt)
+ for cid in np.unique(chunk_lin):
+ cc = np.unravel_index(int(cid), grid)
+ block = read_input_chunk(int(cid), cc) # ONE input chunk at a time
+ mask = chunk_lin == cid
+ local = tuple(in_coords[d][mask] - cc[d] * in_chunks[d] for d in range(in_ndim))
+ obuf[mask] = block[local]
+ out[o_sl] = obuf.reshape(o_shape)
+
+ return output_path
diff --git a/src/dyna_zarr/tiff_reader.py b/src/dyna_zarr/tiff_reader.py
index fa40926..e33412c 100644
--- a/src/dyna_zarr/tiff_reader.py
+++ b/src/dyna_zarr/tiff_reader.py
@@ -1,314 +1,35 @@
-"""
-TIFF File Reader with Concurrent Access Support
-
-Provides lazy, thread-safe reading of TIFF files using tifffile's zarr bridge.
-This enables efficient parallel I/O for multi-dimensional scientific TIFF files.
+"""Lazy TIFF reading via tifffile's zarr bridge.
+
+A TIFF is opened as a lazy, chunked zarr array through ``tifffile``'s ``aszarr`` bridge and
+handed straight to :class:`~dyna_zarr.DynamicArray`. All laziness, slicing, chunk-wise reads
+and memory-boundedness then come from ``DynamicArray``'s pull model (``SliceTransform`` /
+``_read_direct``) -- exactly as for a zarr store -- so there is no TIFF-specific slice logic
+to maintain. The tifffile store stays open for as long as the returned array references it
+(released on garbage collection), which is what lazy TIFF reading requires.
"""
import tifffile
-import tensorstore as ts
-import numpy as np
-from pathlib import Path
-import json
import zarr
-class TiffZarrReader:
- """
- Lazy reader for TIFF files with concurrent access support.
-
- Uses tifffile's aszarr() method to expose TIFF as a zarr array,
- enabling thread-safe parallel reads. The interface mimics TensorStore
- for compatibility with existing code.
-
- Performance characteristics:
- - Opening: Very fast (~0.02s) - only reads metadata
- - Slicing: Lazy - stores slice info, defers data loading until .result() or .read().result()
- - Concurrent reads: Thread-safe via zarr's locking mechanism
- - Limitation: TIFF single-file format doesn't parallelize as well as
- multi-file Zarr format (use large regions, moderate thread counts)
- """
- def __init__(self, tiff_path, zarr_array=None, slice_key=None, parent_shape=None):
- if zarr_array is None:
- # Initial construction from path
- self.tiff_path = str(Path(tiff_path).resolve())
- self._tif = tifffile.TiffFile(self.tiff_path)
- raw_store = self._tif.aszarr()
-
- # Open with zarr - keeps it lazy
- self._zarr_array = zarr.open(raw_store, mode='r')
- self._slice_key = None # No slicing yet
- self._parent_shape = None
-
- self.shape = tuple(self._zarr_array.shape)
- self.dtype = self._zarr_array.dtype
- self.chunks = self._zarr_array.chunks if hasattr(self._zarr_array, 'chunks') else None
- else:
- # Construction from sliced zarr array - store slice for lazy evaluation
- self.tiff_path = tiff_path
- self._tif = None
- self._zarr_array = zarr_array
- self._slice_key = slice_key # Store the slice, don't apply it yet!
- self._parent_shape = parent_shape # Shape before this slice was applied
-
- # Compute shape from the slice WITHOUT loading data
- if slice_key is not None:
- # Use parent_shape (previous slice's result shape) if available
- base_shape = parent_shape if parent_shape is not None else zarr_array.shape
- self.shape = self._compute_sliced_shape(base_shape, slice_key)
- else:
- self.shape = tuple(parent_shape if parent_shape is not None else zarr_array.shape)
-
- self.dtype = zarr_array.dtype
- self.chunks = None
-
- def _compute_sliced_shape(self, original_shape, key):
- """Compute the shape that would result from slicing, without actually slicing."""
- # Normalize the key to a tuple
- if not isinstance(key, tuple):
- key = (key,)
-
- # Handle chained slices - if key is (first_key, second_key), we need to think differently
- if len(key) == 2 and not isinstance(key[0], (int, slice)):
- # This is a chained slice from _combine_slices - DON'T use it for shape computation
- # Instead, just use the last slice in the chain
- # Actually, this shouldn't happen because parent_shape should already account for earlier slices
- pass
-
- # Pad with full slices if needed
- key = key + (slice(None),) * (len(original_shape) - len(key))
-
- new_shape = []
- for dim_size, idx in zip(original_shape, key):
- if isinstance(idx, slice):
- start, stop, step = idx.indices(dim_size)
- length = len(range(start, stop, step))
- new_shape.append(length)
- elif isinstance(idx, int):
- # Integer indexing removes dimension
- continue
- else:
- # For other types (arrays, etc), just use original size
- new_shape.append(dim_size)
-
- return tuple(new_shape)
-
- def __getitem__(self, key):
- """
- Lazy slicing - stores slice info without loading data.
-
- IMPORTANT: This does NOT load data! It returns a new TiffZarrReader
- with the slice stored for later evaluation when .result() is called.
- """
- # For shape computation, we only need the NEW key applied to current shape
- # For materialization, we need ALL keys applied sequentially
- # So store both: combined_key for materialization, key for shape computation
-
- if self._slice_key is not None:
- # Chain slices together for materialization
- combined_key = self._combine_slices(self._slice_key, key)
- else:
- combined_key = key
-
- # Compute shape using ONLY the new key on current shape
- # (parent_shape is passed as self.shape, which already includes previous slices)
- new_shape = self._compute_sliced_shape(self.shape, key)
-
- # Create new reader
- new_reader = TiffZarrReader(
- self.tiff_path,
- zarr_array=self._zarr_array,
- slice_key=combined_key,
- parent_shape=new_shape # Pass computed shape as parent for next slice
- )
- # Override the shape with our computed one
- new_reader.shape = new_shape
-
- return new_reader
-
- def _combine_slices(self, first_key, second_key):
- """Combine two slice operations into one."""
- # For simplicity, we can just store them as a tuple
- # and apply them sequentially when needed
- # A more sophisticated implementation could optimize this
- return (first_key, second_key)
-
- def read(self):
- """
- Async-style read that returns a Future-like object.
-
- For compatibility with TensorStore-style code that uses .read().result()
- This is when data is ACTUALLY loaded from disk.
- """
- class FutureResult:
- def __init__(self, reader):
- self._reader = reader
- def result(self):
- return self._reader._materialize()
-
- return FutureResult(self)
-
- def result(self):
- """Direct result() call - THIS is when data is actually loaded from disk."""
- return self._materialize()
-
- def _materialize(self):
- """
- Actually load the data from disk by applying stored slices.
-
- IMPORTANT: This is where zarr array slicing happens, which materializes data!
- Optimized to apply slices sequentially without loading full array.
- """
- if self._slice_key is None:
- # No slicing, return full array
- return np.asarray(self._zarr_array[:])
- else:
- # Apply the stored slice(s)
- # Check if this is a chained slice (stored as tuple by _combine_slices)
- if self._is_chained_slice(self._slice_key):
- # Chained slices - apply sequentially on zarr array, then on numpy arrays
- slices_to_apply = self._flatten_chained_slices(self._slice_key)
-
- # Apply first slice on zarr array (this loads data from disk)
- result = self._zarr_array[slices_to_apply[0]]
- if not isinstance(result, np.ndarray):
- result = np.asarray(result)
-
- # Apply remaining slices on numpy array (in-memory)
- for s in slices_to_apply[1:]:
- result = result[s]
-
- return result
- else:
- # Single slice
- result = self._zarr_array[self._slice_key]
- if isinstance(result, np.ndarray):
- return result
- else:
- # Scalar or other type
- return np.asarray(result)
-
- def _is_chained_slice(self, key):
- """Check if a key represents chained slices from _combine_slices."""
- # A chained slice is a 2-tuple where first element is not a slice/int
- if isinstance(key, tuple) and len(key) == 2:
- first = key[0]
- # If first element is itself a tuple or a chained slice, it's chained
- if isinstance(first, tuple):
- return True
- return False
-
- def _flatten_chained_slices(self, key):
- """Flatten nested chained slices into a list of slice tuples."""
- slices = []
- if self._is_chained_slice(key):
- first_key, second_key = key
- # Recursively flatten first_key if it's also chained
- if self._is_chained_slice(first_key):
- slices.extend(self._flatten_chained_slices(first_key))
- else:
- slices.append(first_key)
- # Add second_key
- slices.append(second_key)
- else:
- slices.append(key)
- return slices
-
- def __array__(self, dtype=None):
- """
- Support numpy array protocol for np.asarray() conversion.
-
- CRITICAL: This MUST call _materialize() to respect stored slices!
- Otherwise np.asarray() would load the entire array, ignoring slices.
- """
- arr = self._materialize()
- if dtype is not None:
- return arr.astype(dtype)
- return arr
-
- @property
- def spec(self):
- """Provide a spec property for compatibility with TensorStore-style code."""
- return {
- 'driver': 'tiff+zarr',
- 'path': self.tiff_path,
- 'shape': self.shape,
- 'dtype': str(self.dtype),
- 'chunks': self.chunks
- }
-
-
-class ArrayResult:
- """Wrapper for already-materialized numpy arrays."""
- def __init__(self, array):
- self._array = array
- self.shape = array.shape
- self.dtype = array.dtype
-
- def read(self):
- class FutureResult:
- def __init__(self, data):
- self._data = data
- def result(self):
- return self._data
- return FutureResult(self._array)
-
- def result(self):
- """Direct result() call."""
- return self._array
-
- def __array__(self, dtype=None):
- """Support numpy array protocol."""
- if dtype is not None:
- return self._array.astype(dtype)
- return self._array
-
- def __getitem__(self, key):
- return ArrayResult(self._array[key])
-
-
-class ScalarResult:
- """Wrapper for scalar values."""
- def __init__(self, value):
- self._value = value
- self.shape = ()
- self.dtype = np.array(value).dtype
-
- def read(self):
- class FutureResult:
- def __init__(self, data):
- self._data = data
- def result(self):
- return self._data
- return FutureResult(self._value)
-
- def result(self):
- """Direct result() call."""
- return self._value
-
- def __array__(self, dtype=None):
- """Support numpy array protocol."""
- arr = np.array(self._value)
- if dtype is not None:
- return arr.astype(dtype)
- return arr
+def open_tiff_zarr(path):
+ """Open a TIFF as a lazy zarr array via tifffile's zarr bridge (the raw backend object).
-
-def read_tiff_lazy(tiff_path):
- """
- Open a TIFF file for lazy, thread-safe reading.
-
- Fast opening (~0.02s) - only reads metadata, not pixel data.
- Returns a reader that provides:
- - Lazy slicing: data only loaded on access
- - Thread-safe: multiple threads can read concurrently
- - TensorStore-compatible API: supports .read().result() pattern
- - Efficient chunked access via tifffile's zarr bridge
-
- Args:
- tiff_path: Path to the TIFF file
-
- Returns:
- TiffZarrReader with lazy, concurrent read support
+ For a plain single-series TIFF this is a ``zarr.Array``; for a multi-series/multi-level
+ file tifffile yields a group, in which case the first array (series/level 0) is returned.
"""
- return TiffZarrReader(tiff_path)
+ store = tifffile.imread(str(path), aszarr=True) # ZarrTiffStore (lazy; reads on access)
+ obj = zarr.open(store, mode="r")
+ if isinstance(obj, zarr.Group):
+ for key in obj: # multi-series/level -> take the first array
+ item = obj[key]
+ if isinstance(item, zarr.Array):
+ return item
+ raise ValueError(f"TIFF at {path!r} exposes no readable array via aszarr")
+ return obj
+
+
+def read_tiff_lazy(path):
+ """Read a TIFF as a lazy, memory-bounded :class:`~dyna_zarr.DynamicArray`."""
+ from .dynamic_array import DynamicArray
+ return DynamicArray(open_tiff_zarr(path))
diff --git a/tests/test_creation.py b/tests/test_creation.py
new file mode 100644
index 0000000..c8c8dd0
--- /dev/null
+++ b/tests/test_creation.py
@@ -0,0 +1,79 @@
+"""
+Creation ops (nullary generative sources): zeros/ones/full/empty/random + *_like.
+
+These synthesize each region lazily (no underlying array), compose with other ops, and
+stream to disk. ``random`` must be position-deterministic: any sub-slice equals the whole
+array sliced, and io.write matches compute.
+"""
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, io, operations as ops
+
+
+def random_key(rng, shape):
+ key = []
+ for size in shape:
+ c = rng.integers(0, 4)
+ if c == 0:
+ key.append(int(rng.integers(0, size)))
+ elif c == 1:
+ key.append(slice(None))
+ elif c == 2:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b)))
+ else:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b), int(rng.integers(1, 3))))
+ return tuple(key)
+
+
+def test_deterministic_creation():
+ for op, npf in [(ops.zeros, np.zeros), (ops.ones, np.ones)]:
+ a = op((4, 6), dtype=np.float32)
+ assert a.shape == (4, 6) and a.dtype == np.float32
+ np.testing.assert_array_equal(a.compute(), npf((4, 6), np.float32))
+ np.testing.assert_array_equal(ops.full((3, 5), 7.0).compute(), np.full((3, 5), 7.0))
+ # full infers dtype from fill_value
+ assert ops.full((2, 2), 3).dtype == np.array(3).dtype
+
+
+def test_creation_subslice_and_chain():
+ z = ops.ones((8, 8), dtype=np.float32)
+ np.testing.assert_array_equal(z[2:5, ::2].compute(), np.ones((8, 8), np.float32)[2:5, ::2])
+ np.testing.assert_allclose((ops.ones((4, 4)) * 3 + 1).compute(), np.full((4, 4), 4.0))
+
+
+def test_like():
+ base = ops.zeros((3, 5), dtype=np.int16)
+ assert ops.full_like(base, 9).dtype == np.int16
+ np.testing.assert_array_equal(ops.full_like(base, 9).compute(), np.full((3, 5), 9, np.int16))
+ np.testing.assert_array_equal(ops.ones_like(base).compute(), np.ones((3, 5), np.int16))
+
+
+def test_random_position_deterministic():
+ r = ops.random((6, 8, 10), seed=42)
+ whole = r.compute()
+ assert whole.shape == (6, 8, 10) and whole.dtype == np.float32
+ assert 0.0 <= whole.min() and whole.max() < 1.0
+ np.testing.assert_array_equal(whole, r.compute()) # idempotent
+ rng = np.random.default_rng(1)
+ for _ in range(40): # chunk-invariant sub-slices
+ k = random_key(rng, (6, 8, 10))
+ np.testing.assert_array_equal(np.asarray(r[k].compute()), whole[k])
+ # separate seeds differ
+ assert not np.array_equal(ops.random((4, 4), seed=1).compute(),
+ ops.random((4, 4), seed=2).compute())
+
+
+def test_write_generative(tmp_path):
+ for name, pipe, ref in [
+ ("full", ops.full((16, 32, 32), 5.0, dtype=np.float32), np.full((16, 32, 32), 5.0, np.float32)),
+ ("random", ops.random((16, 32, 32), seed=7), None),
+ ]:
+ out = str(tmp_path / f"{name}.zarr")
+ io.write(pipe, out, zarr_format=2, chunks=(4, 16, 16))
+ written = zarr.open(out, mode="r")[:]
+ expected = ref if ref is not None else pipe.compute() # write must equal compute
+ np.testing.assert_array_equal(written, expected)
diff --git a/tests/test_neighborhood.py b/tests/test_neighborhood.py
new file mode 100644
index 0000000..6a353ef
--- /dev/null
+++ b/tests/test_neighborhood.py
@@ -0,0 +1,122 @@
+"""
+Correctness + chunk-invariance for the map_overlap neighbourhood primitive.
+
+Each neighbourhood filter must equal scipy applied to the *whole* array, no matter what
+region size drives the read. We check two things against a scipy reference:
+
+1. **Full compute** ``op(da).compute() == scipy(arr)``
+2. **Random sub-slice** ``op(da)[k] == scipy(arr)[k]`` over many random keys AND several
+ chunkings -- the real test that each read pulls its own halo correctly (chunk/region
+ invariance) and that boundaries are handled exactly.
+"""
+
+import numpy as np
+import pytest
+import zarr
+from scipy import ndimage as ndi
+
+from dyna_zarr import DynamicArray, operations as ops
+
+
+SHAPE = (12, 16, 20)
+
+
+def da_from(arr, chunks):
+ return DynamicArray(zarr.array(arr, chunks=chunks))
+
+
+def random_key(rng, shape):
+ key = []
+ for size in shape:
+ c = rng.integers(0, 4)
+ if c == 0:
+ key.append(int(rng.integers(0, size)))
+ elif c == 1:
+ key.append(slice(None))
+ elif c == 2:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b)))
+ else:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b), int(rng.integers(1, 3))))
+ return tuple(key)
+
+
+_KERNEL = np.random.default_rng(7).random((3, 3, 3)).astype(np.float32) # fixed conv kernel
+
+
+# name -> (dyna_op, scipy_reference); reference uses mode='reflect' to match default boundary
+OPS = {
+ "convolve_3": (lambda da: ops.convolve(da, _KERNEL),
+ lambda a: ndi.convolve(a, _KERNEL, mode="reflect")),
+ "correlate_3": (lambda da: ops.correlate(da, _KERNEL),
+ lambda a: ndi.correlate(a, _KERNEL, mode="reflect")),
+ "gaussian_s2": (lambda da: ops.gaussian_filter(da, 2.0),
+ lambda a: ndi.gaussian_filter(a, 2.0, mode="reflect")),
+ "gaussian_aniso": (lambda da: ops.gaussian_filter(da, (1.0, 2.0, 0.5)),
+ lambda a: ndi.gaussian_filter(a, (1.0, 2.0, 0.5), mode="reflect")),
+ "uniform_5": (lambda da: ops.uniform_filter(da, 5),
+ lambda a: ndi.uniform_filter(a, 5, mode="reflect")),
+ "uniform_even_4": (lambda da: ops.uniform_filter(da, 4),
+ lambda a: ndi.uniform_filter(a, 4, mode="reflect")),
+ "median_3": (lambda da: ops.median_filter(da, 3),
+ lambda a: ndi.median_filter(a, 3, mode="reflect")),
+ "minimum_3": (lambda da: ops.minimum_filter(da, 3),
+ lambda a: ndi.minimum_filter(a, 3, mode="reflect")),
+ "maximum_3": (lambda da: ops.maximum_filter(da, 3),
+ lambda a: ndi.maximum_filter(a, 3, mode="reflect")),
+ "grey_erosion_3": (lambda da: ops.grey_erosion(da, 3),
+ lambda a: ndi.grey_erosion(a, size=3, mode="reflect")),
+ "grey_dilation_3": (lambda da: ops.grey_dilation(da, 3),
+ lambda a: ndi.grey_dilation(a, size=3, mode="reflect")),
+ "laplace": (lambda da: ops.laplace(da),
+ lambda a: ndi.laplace(a, mode="reflect")),
+ "gaussian_laplace": (lambda da: ops.gaussian_laplace(da, 2.0),
+ lambda a: ndi.gaussian_laplace(a, 2.0, mode="reflect")),
+ "gauss_grad_mag": (lambda da: ops.gaussian_gradient_magnitude(da, 2.0),
+ lambda a: ndi.gaussian_gradient_magnitude(a, 2.0, mode="reflect")),
+}
+
+
+@pytest.fixture(scope="module")
+def arr():
+ return np.random.default_rng(0).random(SHAPE).astype(np.float32)
+
+
+@pytest.mark.parametrize("op_name", list(OPS))
+def test_full_compute_matches_scipy(arr, op_name):
+ dyna_op, ref_op = OPS[op_name]
+ da = da_from(arr, chunks=(4, 4, 5))
+ np.testing.assert_allclose(dyna_op(da).compute(), ref_op(arr), atol=1e-5,
+ err_msg=f"{op_name}: full compute != scipy")
+
+
+@pytest.mark.parametrize("chunks", [(12, 16, 20), (3, 4, 5), (5, 7, 6), (1, 16, 20)])
+def test_random_subslice_chunk_invariant(arr, chunks):
+ """op(da)[k] == scipy(arr)[k] over random keys and several chunkings."""
+ rng = np.random.default_rng(1234)
+ failures = []
+ for op_name, (dyna_op, ref_op) in OPS.items():
+ da = da_from(arr, chunks=chunks)
+ ref = ref_op(arr)
+ result = dyna_op(da)
+ for _ in range(15):
+ k = random_key(rng, SHAPE)
+ try:
+ got = result[k].compute()
+ np.testing.assert_allclose(got, ref[k], atol=1e-5)
+ except Exception as e:
+ failures.append(f"{op_name} chunks={chunks} key={k}: "
+ f"{type(e).__name__}: {str(e).splitlines()[-1][:80]}")
+ break
+ assert not failures, "map_overlap chunk-invariance failures:\n" + "\n".join(failures)
+
+
+def test_write_readback_matches_scipy(arr, tmp_path):
+ """End-to-end: gaussian through io.write (region-streamed) == scipy(whole)."""
+ from dyna_zarr import io
+ da = da_from(arr, chunks=(4, 4, 5))
+ out = str(tmp_path / "gauss.zarr")
+ io.write(ops.gaussian_filter(da, 2.0), out, zarr_format=2, chunks=(4, 4, 5))
+ written = zarr.open(out, mode="r")[:]
+ np.testing.assert_allclose(written, ndi.gaussian_filter(arr, 2.0, mode="reflect"), atol=1e-5)
diff --git a/tests/test_new_ops.py b/tests/test_new_ops.py
new file mode 100644
index 0000000..9626c9c
--- /dev/null
+++ b/tests/test_new_ops.py
@@ -0,0 +1,94 @@
+"""
+Correctness + chunk-invariance for the ops added to match ome_zarr_pro's dask surface:
+Tier 1 (isin, digitize, rot90) and Tier 2 (var, std, argmin, argmax, diff, gradient).
+
+Each must equal numpy for full compute AND for random sub-slices across several chunkings
+(the read-key math / streaming must be region-independent).
+"""
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, operations as ops
+
+SHAPE = (6, 8, 10)
+
+
+def da_from(arr, chunks):
+ return DynamicArray(zarr.array(arr, chunks=chunks))
+
+
+def random_key(rng, shape):
+ key = []
+ for size in shape:
+ c = rng.integers(0, 4)
+ if c == 0:
+ key.append(int(rng.integers(0, size)))
+ elif c == 1:
+ key.append(slice(None))
+ elif c == 2:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b)))
+ else:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b), int(rng.integers(1, 3))))
+ return tuple(key)
+
+
+@pytest.fixture(scope="module")
+def arr():
+ rng = np.random.default_rng(0)
+ a = rng.random(SHAPE).astype(np.float32)
+ a[a < 0.1] = 0.25 # a few exact bin values for isin/digitize
+ return a
+
+
+# name -> (dyna_op, numpy_ref, exact?) -- all support sub-slicing
+OPS = {
+ "isin": (lambda d: ops.isin(d, [0.25, 0.5]), lambda a: np.isin(a, [0.25, 0.5]), True),
+ "digitize": (lambda d: ops.digitize(d, [0.25, 0.5, 0.75]), lambda a: np.digitize(a, [0.25, 0.5, 0.75]), True),
+ "rot90_k1": (lambda d: ops.rot90(d, 1, (1, 2)), lambda a: np.rot90(a, 1, (1, 2)), False),
+ "rot90_k3": (lambda d: ops.rot90(d, 3, (0, 2)), lambda a: np.rot90(a, 3, (0, 2)), False),
+ "var_a0": (lambda d: ops.var(d, 0), lambda a: a.var(0), False),
+ "std_a1": (lambda d: ops.std(d, 1), lambda a: a.std(1), False),
+ "var_ddof1": (lambda d: ops.var(d, 2, ddof=1), lambda a: a.var(2, ddof=1), False),
+ "argmin_a0": (lambda d: ops.argmin(d, 0), lambda a: a.argmin(0), True),
+ "argmax_a2": (lambda d: ops.argmax(d, 2), lambda a: a.argmax(2), True),
+ "diff_a2": (lambda d: ops.diff(d, axis=2), lambda a: np.diff(a, axis=2), False),
+ "diff_n2_a0": (lambda d: ops.diff(d, 2, axis=0), lambda a: np.diff(a, 2, axis=0), False),
+ "gradient_a1":(lambda d: ops.gradient(d, axis=1), lambda a: np.gradient(a, axis=1), False),
+}
+
+
+@pytest.mark.parametrize("name", list(OPS))
+def test_full_compute_matches_numpy(arr, name):
+ dyna_op, ref_op, exact = OPS[name]
+ got = np.asarray(dyna_op(da_from(arr, (2, 3, 4))).compute())
+ ref = np.asarray(ref_op(arr))
+ assert got.shape == ref.shape, f"{name}: shape {got.shape} != {ref.shape}"
+ if exact:
+ np.testing.assert_array_equal(got, ref)
+ else:
+ np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("chunks", [SHAPE, (2, 3, 4), (1, 8, 10), (3, 1, 5)])
+def test_random_subslice_chunk_invariant(arr, chunks):
+ rng = np.random.default_rng(7)
+ failures = []
+ for name, (dyna_op, ref_op, exact) in OPS.items():
+ result = dyna_op(da_from(arr, chunks))
+ ref = np.asarray(ref_op(arr))
+ for _ in range(12):
+ k = random_key(rng, ref.shape)
+ try:
+ got = np.asarray(result[k].compute())
+ if exact:
+ np.testing.assert_array_equal(got, ref[k])
+ else:
+ np.testing.assert_allclose(got, ref[k], rtol=1e-4, atol=1e-4)
+ except Exception as e:
+ failures.append(f"{name} chunks={chunks} key={k}: "
+ f"{type(e).__name__}: {str(e).splitlines()[-1][:70]}")
+ break
+ assert not failures, "new-op chunk-invariance failures:\n" + "\n".join(failures)
diff --git a/tests/test_numpy_protocol.py b/tests/test_numpy_protocol.py
new file mode 100644
index 0000000..0fbd9cb
--- /dev/null
+++ b/tests/test_numpy_protocol.py
@@ -0,0 +1,130 @@
+"""
+NumPy-protocol surface on DynamicArray: array METHODS (.astype/.clip/.round) and the ufunc
+protocol (np.sqrt(a), np.add(a, 2), np.clip(a, lo, hi)). These make a DynamicArray a drop-in
+for numpy/dask arrays (used by the ome_zarr_pro backend). Distinct from operations.* function
+fuzzing (test_transform_correctness) -- e.g. np.clip goes through numpy's _wrapfunc calling
+a.clip(min, max, out=...), NOT the ufunc protocol, which once silently returned wrong data.
+"""
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray
+
+
+@pytest.fixture
+def da_arr():
+ arr = (np.random.default_rng(0).random((3, 8, 8)).astype(np.float32) * 10) - 5
+ return DynamicArray(zarr.array(arr, chunks=(1, 4, 4))), arr # multi-chunk -> block reads
+
+
+# --- array methods ---
+
+def test_astype_method(da_arr):
+ d, arr = da_arr
+ np.testing.assert_array_equal(d.astype("int16").compute(), arr.astype("int16"))
+
+
+def test_clip_method(da_arr):
+ d, arr = da_arr
+ np.testing.assert_allclose(d.clip(-1, 2).compute(), np.clip(arr, -1, 2))
+ np.testing.assert_allclose(d.clip(0, None).compute(), np.clip(arr, 0, None))
+
+
+def test_round_method(da_arr):
+ d, arr = da_arr
+ np.testing.assert_allclose(d.round().compute(), np.round(arr))
+ np.testing.assert_allclose(d.round(1).compute(), np.round(arr, 1))
+
+
+def test_clip_round_reject_out(da_arr):
+ d, _ = da_arr
+ with pytest.raises(TypeError):
+ d.clip(0, 1, out=np.empty(d.shape, dtype=d.dtype))
+ with pytest.raises(TypeError):
+ d.round(0, out=np.empty(d.shape, dtype=d.dtype))
+
+
+# --- numpy top-level functions that dispatch onto the array (the np.clip _wrapfunc path) ---
+
+def test_np_clip_and_round(da_arr):
+ d, arr = da_arr
+ np.testing.assert_allclose(np.clip(d, -1, 2).compute(), np.clip(arr, -1, 2)) # block-wise, was buggy
+ np.testing.assert_allclose(np.round(d).compute(), np.round(arr))
+
+
+# --- ufunc protocol (__array_ufunc__) ---
+
+@pytest.mark.parametrize("npf,ref", [
+ (np.sqrt, lambda a: np.sqrt(np.abs(a))),
+ (np.absolute, np.abs), # alias: absolute -> abs
+ (np.abs, np.abs),
+ (np.negative, np.negative),
+ (np.exp, lambda a: np.exp(np.clip(a, None, 5))),
+ (np.sign, np.sign),
+ (np.floor, np.floor),
+ (np.ceil, np.ceil),
+])
+def test_unary_ufunc_dispatch(da_arr, npf, ref):
+ d, arr = da_arr
+ # guard domain for sqrt/exp via the ref's own transform on d
+ got = npf(np.abs(d)).compute() if npf is np.sqrt else (
+ npf(np.clip(d, None, 5)).compute() if npf is np.exp else npf(d).compute())
+ np.testing.assert_allclose(got, ref(arr), rtol=1e-5)
+
+
+def test_binary_ufunc_dispatch(da_arr):
+ d, arr = da_arr
+ np.testing.assert_allclose(np.add(d, 2).compute(), arr + 2)
+ np.testing.assert_allclose(np.add(2, d).compute(), 2 + arr) # scalar on left
+ np.testing.assert_allclose(np.multiply(d, d).compute(), arr * arr)
+ np.testing.assert_array_equal(np.greater(d, 0).compute(), arr > 0)
+ np.testing.assert_allclose(np.true_divide(d, 3).compute(), arr / 3) # alias -> divide
+
+
+def test_unsupported_ufunc_errors_cleanly(da_arr):
+ """A ufunc dyna_zarr doesn't implement must NOT silently return wrong data."""
+ d, _ = da_arr
+ with pytest.raises(TypeError):
+ np.sin(d).compute()
+
+
+# --- operators (dunders) round-trip, incl. chaining ---
+
+def test_operators_and_chain(da_arr):
+ d, arr = da_arr
+ np.testing.assert_array_equal(((d > -1) & (d < 2)).compute(), (arr > -1) & (arr < 2))
+ np.testing.assert_allclose((-d).compute(), -arr)
+ np.testing.assert_allclose(((d.astype("float32") + 5) / 2).clip(0, 4).compute(),
+ np.clip((arr + 5) / 2, 0, 4))
+
+
+# --- dask-compat methods used by the ome_zarr_pro backend ---
+
+def test_rechunk_persist_are_noops(da_arr):
+ d, arr = da_arr
+ np.testing.assert_array_equal(d.rechunk({0: d.shape[0]}).persist().compute(), arr)
+ np.testing.assert_array_equal(d.rechunk(-1).compute(), arr)
+
+
+def test_map_blocks_method_dask_style(da_arr):
+ d, arr = da_arr
+ # dask-style call: func, dtype, and a dask-only meta kwarg that must be ignored
+ got = d.map_blocks(lambda b: b * 2, dtype="float32", meta=np.array((), dtype="float32"))
+ np.testing.assert_allclose(got.compute(), arr * 2)
+ # a bound (non-dask) kwarg is forwarded to func
+ got2 = d.map_blocks(lambda b, k=1.0: b + k, dtype="float32", k=3.0)
+ np.testing.assert_allclose(got2.compute(), arr + 3.0)
+ with pytest.raises(NotImplementedError):
+ d.map_blocks(lambda b: b, dtype="float32", drop_axis=0)
+
+
+def test_map_overlap_method_dask_style(da_arr):
+ from scipy import ndimage as ndi
+ d, arr = da_arr
+ got = d.map_overlap(lambda b: ndi.uniform_filter(b, 3, mode="reflect"),
+ depth=1, boundary="reflect", trim=True, dtype="float32",
+ meta=np.array((), dtype="float32"))
+ np.testing.assert_allclose(got.compute(), ndi.uniform_filter(arr, 3, mode="reflect"), atol=1e-5)
+ with pytest.raises(NotImplementedError):
+ d.map_overlap(lambda b: b, depth=1, trim=False)
diff --git a/tests/test_rechunk.py b/tests/test_rechunk.py
new file mode 100644
index 0000000..f92a2d9
--- /dev/null
+++ b/tests/test_rechunk.py
@@ -0,0 +1,208 @@
+"""
+Two-phase disk-staged rechunk engine (Rechunker's algorithm, dask-free) + flatten built on
+top of it. Covers all grid-alignment cases (coarsen / refine / misaligned / mixed / partial),
+read-once, and the flatten bridge (rechunk-to-flat-contiguous + relabel).
+"""
+import os
+import threading
+import time
+
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, io, operations as ops
+from dyna_zarr.rechunk import rechunk, flatten_write, reshape_write
+
+
+class CountZarr:
+ """Wrap a zarr array and count element reads (to assert read-once). Thread-safe so it's
+ valid under the parallel copy engine."""
+
+ def __init__(self, z):
+ self.z = z
+ self.reads = 0
+ self._lock = threading.Lock()
+
+ shape = property(lambda self: self.z.shape)
+ chunks = property(lambda self: self.z.chunks)
+ dtype = property(lambda self: self.z.dtype)
+
+ def __getitem__(self, k):
+ b = self.z[k]
+ with self._lock:
+ self.reads += int(np.prod(b.shape))
+ return b
+
+
+@pytest.fixture
+def src3d():
+ arr = np.random.default_rng(0).random((20, 24, 18)).astype(np.float32)
+ return arr, zarr.array(arr, chunks=(8, 8, 8))
+
+
+@pytest.mark.parametrize("tc", [
+ (16, 16, 16), # coarsen (source divides target) -> one-pass consolidate
+ (4, 4, 4), # refine (target divides source) -> one-pass split
+ (5, 7, 6), # misaligned -> two-pass via disk intermediate
+ (8, 6, 18), # mixed: divides on axis0, misaligned axis1, full axis2
+ (20, 24, 18), # single big chunk (partial edges)
+])
+def test_rechunk_matches_source(src3d, tc, tmp_path):
+ arr, z = src3d
+ out = str(tmp_path / "o.zarr")
+ rechunk(z, tc, out, zarr_format=2)
+ r = zarr.open(out, mode="r")
+ np.testing.assert_array_equal(r[:], arr)
+ assert r.chunks == tc
+
+
+def test_rechunk_read_once(src3d, tmp_path):
+ arr, z = src3d
+ cz = CountZarr(z)
+ rechunk(cz, (5, 7, 6), str(tmp_path / "o.zarr"), zarr_format=2) # misaligned = two-pass
+ assert cz.reads == arr.size # each source element read exactly once
+
+
+@pytest.mark.parametrize("max_workers", [1, 4])
+def test_rechunk_parallel_correct_and_read_once(src3d, max_workers, tmp_path):
+ """Small max_mem forces MANY regions -> the thread pool is actually exercised. Result must
+ match numpy and each source element must still be read exactly once under concurrency."""
+ arr, z = src3d
+ cz = CountZarr(z)
+ out = str(tmp_path / "o.zarr")
+ rechunk(cz, (5, 7, 6), out, max_mem=4096, max_workers=max_workers, zarr_format=2) # misaligned
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr)
+ assert cz.reads == arr.size
+
+
+def test_rechunk_dynamic_source(src3d, tmp_path):
+ arr, z = src3d
+ da = DynamicArray(z)
+ out = str(tmp_path / "o.zarr")
+ rechunk(da, (5, 7, 6), out, zarr_format=2)
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr)
+
+
+@pytest.mark.parametrize("shape,chunks", [
+ ((12, 10, 8, 6), (5, 4, 3, 4)), # 4D non-cubic partial
+ ((20, 24, 18), (8, 8, 8)), # 3D cubic
+ ((100,), (16,)), # already 1D
+])
+def test_flatten_write_matches_numpy(shape, chunks, tmp_path):
+ arr = np.random.default_rng(2).random(shape).astype(np.float32)
+ da = DynamicArray(zarr.array(arr, chunks=chunks))
+ out = str(tmp_path / "f.zarr")
+ flatten_write(da, out, output_chunks=(64,), zarr_format=2)
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.ravel())
+
+
+def test_flatten_write_read_once(tmp_path):
+ arr = np.random.default_rng(3).random((20, 24, 18)).astype(np.float32)
+ cz = CountZarr(zarr.array(arr, chunks=(8, 8, 8)))
+ flatten_write(DynamicArray(cz), str(tmp_path / "f.zarr"),
+ output_chunks=(64,), max_mem=1 << 20, zarr_format=2)
+ assert cz.reads == arr.size # source read exactly once through the rechunk stage
+
+
+@pytest.mark.parametrize("max_mem", [1 << 20, 4096, 512, 64]) # force split axis from 0 -> deep
+def test_flatten_write_tiny_budget_deep_split(max_mem, tmp_path):
+ """Small budgets push the split axis past axis 0; result must still equal ravel()."""
+ arr = np.random.default_rng(5).random((6, 5, 4, 7)).astype(np.float32)
+ da = DynamicArray(zarr.array(arr, chunks=(2, 2, 2, 3)))
+ out = str(tmp_path / "f.zarr")
+ flatten_write(da, out, output_chunks=(9,), max_mem=max_mem, zarr_format=2)
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.ravel())
+
+
+def test_iowrite_flatten_routes_to_rechunk(tmp_path):
+ arr = np.random.default_rng(4).random((12, 10, 8, 6)).astype(np.float32)
+ da = DynamicArray(zarr.array(arr, chunks=(5, 4, 3, 4)))
+ out = str(tmp_path / "f.zarr")
+ io.write(ops.flatten(da), out, zarr_format=2, chunks=(64,))
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.ravel())
+
+
+@pytest.mark.parametrize("shape,chunks,target,tchunks", [
+ ((12, 10, 8, 6), (5, 4, 3, 4), (120, 48), (7, 9)), # 4D -> 2D
+ ((12, 10, 8, 6), (5, 4, 3, 4), (24, 10, 24), (7, 4, 10)), # 4D -> 3D
+ ((12, 10, 8, 6), (5, 4, 3, 4), (6, 20, 8, 6), (4, 7, 5, 4)),# 4D -> 4D
+ ((12, 10, 8, 6), (5, 4, 3, 4), (5760,), (100,)), # 4D -> 1D
+ ((20, 24, 18), (8, 8, 8), (60, 144), (16, 16)), # 3D -> 2D
+ ((100,), (16,), (10, 10), (3, 3)), # 1D -> 2D
+])
+def test_reshape_write_matches_numpy(shape, chunks, target, tchunks, tmp_path):
+ arr = np.random.default_rng(6).random(shape).astype(np.float32)
+ da = DynamicArray(zarr.array(arr, chunks=chunks))
+ out = str(tmp_path / "r.zarr")
+ reshape_write(da, target, out, output_chunks=tchunks, zarr_format=2)
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.reshape(target))
+
+
+@pytest.mark.parametrize("max_mem", [1 << 20, 4096, 256]) # force deep flatten split
+def test_reshape_write_tiny_budget(max_mem, tmp_path):
+ arr = np.random.default_rng(7).random((6, 5, 4, 7)).astype(np.float32)
+ da = DynamicArray(zarr.array(arr, chunks=(2, 2, 2, 3)))
+ out = str(tmp_path / "r.zarr")
+ reshape_write(da, (24, 35), out, output_chunks=(5, 8), max_mem=max_mem, zarr_format=2)
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.reshape(24, 35))
+
+
+def test_iowrite_reshape_routes_to_staged(tmp_path):
+ arr = np.random.default_rng(8).random((12, 10, 8, 6)).astype(np.float32)
+ da = DynamicArray(zarr.array(arr, chunks=(5, 4, 3, 4)))
+ out = str(tmp_path / "r.zarr")
+ io.write(ops.reshape(da, (80, 72)), out, zarr_format=2, chunks=(16, 3))
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.reshape(80, 72))
+
+
+@pytest.mark.skipif(os.environ.get("DYNA_LARGE") != "1",
+ reason="5GB flatten; set DYNA_LARGE=1 to run (slow, needs ~10GB disk)")
+def test_flatten_5gb_5d_memory_bounded(tmp_path):
+ """Real large-data check: 5D float32 ~5GB, chunks (1,1,32,32,32). Flatten via io.write must
+ stay memory-bounded (peak << data size) and be correct (spot-checked, can't hold 5GB in RAM)."""
+ import psutil
+ shape = (10, 4, 256, 256, 512) # 1.34e9 elems * 4B = 5.37 GB
+ chunks = (1, 1, 32, 32, 32)
+ assert all(s % c == 0 for s, c in zip(shape[2:], chunks[2:]))
+ data_gb = np.prod(shape) * 4 / 1e9
+
+ src = str(tmp_path / "src.zarr")
+ z = zarr.open(src, mode="w", shape=shape, chunks=chunks, dtype="float32")
+ rng = np.random.default_rng(0)
+ for t in range(shape[0]): # fill one (z,y,x) volume at a time
+ for c in range(shape[1]):
+ z[t, c] = rng.random((shape[2], shape[3], shape[4]), dtype=np.float32)
+
+ out = str(tmp_path / "flat.zarr")
+ proc = psutil.Process()
+ base = proc.memory_info().rss
+ peak = [base]
+ stop = [False]
+
+ def mon():
+ while not stop[0]:
+ peak[0] = max(peak[0], proc.memory_info().rss)
+ time.sleep(0.01)
+
+ th = threading.Thread(target=mon)
+ th.start()
+ io.write(ops.flatten(DynamicArray(z)), out, zarr_format=2, chunks=(1 << 22,))
+ stop[0] = True
+ th.join()
+
+ peak_gb = (peak[0] - base) / 1e9
+ print(f"\n5GB flatten: data {data_gb:.2f}GB peakDelta {peak_gb:.2f}GB")
+
+ of = zarr.open(out, mode="r")
+ N = int(np.prod(shape))
+ assert of.shape == (N,)
+ # spot-check correctness: random flat indices vs source at unravelled coords
+ idx = np.sort(rng.integers(0, N, size=500))
+ coords = np.unravel_index(idx, shape)
+ expect = np.array([z[tuple(int(coords[d][i]) for d in range(len(shape)))]
+ for i in range(len(idx))], dtype=np.float32)
+ np.testing.assert_array_equal(of.oindex[idx], expect)
+
+ # memory must be bounded well under the data size (budget ~256MB + overheads, not O(5GB))
+ assert peak_gb < 1.5, f"flatten peak {peak_gb:.2f}GB not bounded vs {data_gb:.2f}GB data"
diff --git a/tests/test_reductions.py b/tests/test_reductions.py
new file mode 100644
index 0000000..97f773d
--- /dev/null
+++ b/tests/test_reductions.py
@@ -0,0 +1,153 @@
+"""
+Correctness + chunk-invariance for the streaming reduce primitive.
+
+A reduction must equal numpy no matter the chunking or the region driving the read, and
+stay memory-bound (the reduced axis is streamed in bounded chunks). We check full compute,
+random sub-slices across several chunkings, keepdims, global (axis=None), and an
+end-to-end write of a projection.
+"""
+
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, operations as ops
+
+SHAPE = (6, 8, 10)
+
+
+def da_from(arr, chunks):
+ return DynamicArray(zarr.array(arr, chunks=chunks))
+
+
+def random_key(rng, shape):
+ key = []
+ for size in shape:
+ c = rng.integers(0, 4)
+ if c == 0:
+ key.append(int(rng.integers(0, size)))
+ elif c == 1:
+ key.append(slice(None))
+ elif c == 2:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b)))
+ else:
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b), int(rng.integers(1, 3))))
+ return tuple(key)
+
+
+# name -> (dyna_op, numpy_ref) ; each reduces to a known shape
+REDUCERS = {
+ "sum_a0": (lambda d: ops.sum(d, 0), lambda a: a.sum(0)),
+ "max_a1": (lambda d: ops.max(d, 1), lambda a: a.max(1)),
+ "min_a2": (lambda d: ops.min(d, 2), lambda a: a.min(2)),
+ "mean_a0": (lambda d: ops.mean(d, 0), lambda a: a.mean(0)),
+ "prod_a1": (lambda d: ops.prod(d, 1), lambda a: a.prod(1)),
+ "sum_a02": (lambda d: ops.sum(d, (0, 2)), lambda a: a.sum(axis=(0, 2))),
+ "max_a0_keep": (lambda d: ops.max(d, 0, keepdims=True), lambda a: a.max(0, keepdims=True)),
+ "mean_a12_keep":(lambda d: ops.mean(d, (1, 2), keepdims=True),
+ lambda a: a.mean(axis=(1, 2), keepdims=True)),
+ "median_a0": (lambda d: ops.median(d, 0), lambda a: np.median(a, 0)),
+ "median_a12": (lambda d: ops.median(d, (1, 2)), lambda a: np.median(a, axis=(1, 2))),
+}
+
+
+@pytest.fixture(scope="module")
+def arr():
+ return np.random.default_rng(0).random(SHAPE).astype(np.float32)
+
+
+@pytest.mark.parametrize("name", list(REDUCERS))
+def test_full_compute_matches_numpy(arr, name):
+ dyna_op, ref_op = REDUCERS[name]
+ got = dyna_op(da_from(arr, (2, 3, 4))).compute()
+ np.testing.assert_allclose(got, ref_op(arr), rtol=1e-5, atol=1e-5,
+ err_msg=f"{name}: full compute")
+
+
+@pytest.mark.parametrize("chunks", [(6, 8, 10), (2, 3, 4), (1, 8, 10), (3, 1, 5)])
+def test_random_subslice_chunk_invariant(arr, chunks):
+ rng = np.random.default_rng(7)
+ failures = []
+ for name, (dyna_op, ref_op) in REDUCERS.items():
+ result = dyna_op(da_from(arr, chunks))
+ ref = ref_op(arr)
+ for _ in range(15):
+ k = random_key(rng, ref.shape)
+ try:
+ np.testing.assert_allclose(result[k].compute(), ref[k], rtol=1e-5, atol=1e-5)
+ except Exception as e:
+ failures.append(f"{name} chunks={chunks} key={k}: "
+ f"{type(e).__name__}: {str(e).splitlines()[-1][:80]}")
+ break
+ assert not failures, "reduce chunk-invariance failures:\n" + "\n".join(failures)
+
+
+def test_global_reductions(arr):
+ da = da_from(arr, (2, 3, 4))
+ for name, npf in [("sum", np.sum), ("max", np.max), ("min", np.min), ("mean", np.mean)]:
+ got = getattr(ops, name)(da).compute()
+ assert np.ndim(got) == 0
+ np.testing.assert_allclose(got, npf(arr), rtol=1e-5, atol=1e-5, err_msg=name)
+
+
+def test_any_all(arr):
+ mask = (arr > 0.5)
+ dm = da_from(mask, (2, 3, 4))
+ np.testing.assert_array_equal(ops.any(dm, 0).compute(), mask.any(0))
+ np.testing.assert_array_equal(ops.all(dm, 1).compute(), mask.all(1))
+ assert bool(ops.any(dm).compute()) == bool(mask.any())
+
+
+def test_streaming_is_memory_bounded(arr):
+ """A tiny strip budget forces many reduced-axis chunks but the result is unchanged
+ (proves the associative streaming path, independent of strip size)."""
+ from dyna_zarr.operations.reductions import ReduceTransform
+ da = da_from(arr, (2, 3, 4))
+ t = ReduceTransform(da, "sum", axis=0, strip_bytes=64) # forces chunk_len == 1
+ streamed = da._with_transform(t).compute()
+ np.testing.assert_allclose(streamed, arr.sum(0), rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize("chunks", [(6, 8, 10), (2, 3, 4), (1, 8, 10)])
+def test_histogram_matches_numpy(arr, chunks):
+ """Streaming histogram == numpy exactly, across chunkings (associativity) and for
+ auto-range, fixed-range, and explicit-edge bins."""
+ da = da_from((arr * 255).astype(np.float32), chunks)
+ a = (arr * 255).astype(np.float32)
+ # auto range
+ c, e = ops.histogram(da, bins=64)
+ cn, en = np.histogram(a, bins=64)
+ np.testing.assert_array_equal(c, cn)
+ np.testing.assert_allclose(e, en)
+ assert c.sum() == a.size and c.dtype == np.int64
+ # fixed range
+ c2, _ = ops.histogram(da, bins=32, range=(0, 255))
+ np.testing.assert_array_equal(c2, np.histogram(a, bins=32, range=(0, 255))[0])
+ # explicit edges + per-plane slice
+ edges = np.linspace(0, 255, 40)
+ np.testing.assert_array_equal(ops.histogram(da, bins=edges)[0], np.histogram(a, bins=edges)[0])
+ np.testing.assert_array_equal(da[2].histogram(64, (0, 255))[0],
+ np.histogram(a[2], bins=64, range=(0, 255))[0])
+
+
+def test_write_projection(arr, tmp_path):
+ """End-to-end: max-projection over axis 0 through io.write == numpy."""
+ from dyna_zarr import io
+ da = da_from(arr, (2, 3, 4))
+ out = str(tmp_path / "proj.zarr")
+ io.write(ops.max(da, 0), out, zarr_format=2, chunks=(4, 5))
+ np.testing.assert_allclose(zarr.open(out, mode="r")[:], arr.max(0), rtol=1e-5, atol=1e-5)
+
+
+def test_unique_matches_numpy(tmp_path):
+ """Streaming unique over the whole array, sorted, like numpy.unique (values only)."""
+ rng = np.random.default_rng(3)
+ lab = rng.integers(0, 12, size=(6, 20, 24)).astype("int32")
+ da = da_from(lab, (2, 7, 5)) # awkward chunks -> multi-region stream
+ np.testing.assert_array_equal(ops.unique(da), np.unique(lab))
+ # float values + empty
+ f = rng.random((4, 8)).astype("float32")
+ np.testing.assert_array_equal(ops.unique(da_from(f, (1, 3))), np.unique(f))
+ assert ops.unique(da_from(np.zeros((0,), "int32"), (1,))).size == 0
diff --git a/tests/test_reindex.py b/tests/test_reindex.py
new file mode 100644
index 0000000..0faf070
--- /dev/null
+++ b/tests/test_reindex.py
@@ -0,0 +1,58 @@
+"""
+Streaming reindex writer for reshape/flatten: output-chunk driven, one input chunk at a
+time. Correct on any chunking (incl. non-cubic + partial chunks), hard-bounded memory,
+race-free. io.write auto-routes reshape/flatten (outermost) to it.
+"""
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, io, operations as ops
+from dyna_zarr.reindex import reindex_write
+
+
+@pytest.fixture
+def da_arr():
+ arr = np.random.default_rng(0).random((12, 10, 8, 6)).astype(np.float32)
+ return DynamicArray(zarr.array(arr, chunks=(5, 4, 3, 4))), arr # non-cubic, partial
+
+
+@pytest.mark.parametrize("out_shape,out_chunks", [
+ ((120, 48), (7, 9)), # 4D -> 2D
+ ((24, 10, 24), (7, 4, 10)), # 4D -> 3D
+ ((6, 20, 8, 6), (4, 7, 5, 4)), # 4D -> 4D
+ ((12 * 10 * 8 * 6,), (100,)), # flatten
+])
+def test_reindex_write_matches_numpy(da_arr, out_shape, out_chunks, tmp_path):
+ da, arr = da_arr
+ out = str(tmp_path / "o.zarr")
+ reindex_write(da, out_shape, out, out_chunks)
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], arr.reshape(out_shape))
+
+
+def test_lru_cache_same_result(da_arr, tmp_path):
+ da, arr = da_arr
+ a = str(tmp_path / "a.zarr")
+ b = str(tmp_path / "b.zarr")
+ reindex_write(da, (120, 48), a, (7, 9), cache_size=0)
+ reindex_write(da, (120, 48), b, (7, 9), cache_size=8)
+ np.testing.assert_array_equal(zarr.open(a, mode="r")[:], zarr.open(b, mode="r")[:])
+ np.testing.assert_array_equal(zarr.open(a, mode="r")[:], arr.reshape(120, 48))
+
+
+def test_iowrite_routes_reshape_flatten(da_arr, tmp_path):
+ da, arr = da_arr
+ r = str(tmp_path / "r.zarr")
+ f = str(tmp_path / "f.zarr")
+ io.write(ops.reshape(da, (80, 72)), r, zarr_format=2, chunks=(16, 3))
+ io.write(ops.flatten(da), f, zarr_format=2, chunks=(64,))
+ np.testing.assert_array_equal(zarr.open(r, mode="r")[:], arr.reshape(80, 72))
+ np.testing.assert_array_equal(zarr.open(f, mode="r")[:], arr.ravel())
+
+
+def test_iowrite_nested_reshape_uses_normal_path(da_arr, tmp_path):
+ """reshape NOT outermost (abs on top) -> normal region writer, still correct."""
+ da, arr = da_arr
+ out = str(tmp_path / "ra.zarr")
+ io.write(ops.abs(ops.reshape(da, (80, 72))), out, zarr_format=2, chunks=(16, 3))
+ np.testing.assert_array_equal(zarr.open(out, mode="r")[:], np.abs(arr).reshape(80, 72))
diff --git a/tests/test_scan.py b/tests/test_scan.py
new file mode 100644
index 0000000..744f8bd
--- /dev/null
+++ b/tests/test_scan.py
@@ -0,0 +1,80 @@
+"""
+Prefix-scan ops (cumsum/cumprod/cummax/cummin). A scan is not chunk-local: output[..p..]
+depends on the whole prefix along the scan axis, so read(key) reads [0, stop) on that axis.
+The sub-slice tests are the real check that the prefix math + squeeze are correct and
+chunk-invariant.
+"""
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, operations as ops
+
+SHAPE = (5, 7, 6)
+REF = {
+ "cumsum": lambda a, ax: np.cumsum(a, axis=ax),
+ "cumprod": lambda a, ax: np.cumprod(a, axis=ax),
+ "cummax": lambda a, ax: np.maximum.accumulate(a, axis=ax),
+ "cummin": lambda a, ax: np.minimum.accumulate(a, axis=ax),
+}
+OP = {"cumsum": ops.cumsum, "cumprod": ops.cumprod, "cummax": ops.cummax, "cummin": ops.cummin}
+
+
+@pytest.fixture
+def data():
+ arr = (np.random.default_rng(0).random(SHAPE).astype(np.float32) * 4) - 2
+ return arr, DynamicArray(zarr.array(arr, chunks=(2, 3, 4)))
+
+
+@pytest.mark.parametrize("name", list(OP))
+@pytest.mark.parametrize("axis", [0, 1, 2, -1])
+def test_full_compute_matches_numpy(data, name, axis):
+ arr, d = data
+ np.testing.assert_allclose(OP[name](d, axis).compute(), REF[name](arr, axis % 3), atol=1e-4)
+
+
+@pytest.mark.parametrize("name", list(OP))
+@pytest.mark.parametrize("axis", [0, 1, 2])
+@pytest.mark.parametrize("key", [
+ np.s_[1:4, 2:6, :], np.s_[2, :, 1:5], np.s_[::2, 1:6:2, 3], np.s_[:, 0, :], np.s_[4, 6, 5],
+])
+def test_subslice_chunk_invariant(data, name, axis, key):
+ """op(d, axis)[key] == numpy(arr, axis)[key] -- prefix read + squeeze, any chunking."""
+ arr, d = data
+ np.testing.assert_allclose(OP[name](d, axis)[key].compute(), REF[name](arr, axis)[key], atol=1e-4)
+
+
+def test_int_dtype_matches_numpy(data):
+ arr, _ = data
+ di = DynamicArray(zarr.array(arr.astype("int16"), chunks=(2, 3, 4)))
+ assert ops.cumsum(di, 0).dtype == np.cumsum(arr.astype("int16"), axis=0).dtype
+
+
+@pytest.mark.parametrize("name", list(OP))
+@pytest.mark.parametrize("axis", [0, 1, 2])
+def test_scan_write_streams_correctly(data, name, axis, tmp_path):
+ """io.write routes an outermost scan to the bounded-carry streaming writer; a tiny
+ region budget forces multiple cross-section tiles AND multiple scan-axis strips."""
+ from dyna_zarr import io
+ arr, d = data
+ out = str(tmp_path / f"{name}{axis}.zarr")
+ io.write(OP[name](d, axis), out, zarr_format=2, chunks=(2, 3, 4), region_size_mb=0.001)
+ np.testing.assert_allclose(zarr.open(out, mode="r")[:], REF[name](arr, axis), atol=1e-4)
+
+
+def test_scan_write_multi_strip_matches(tmp_path):
+ """A tiny budget forces many cross-section tiles AND many scan-axis strips, so the running
+ carry is exercised across strip boundaries (where a naive per-block scan would be wrong)."""
+ from dyna_zarr import io
+ rng = np.random.default_rng(1)
+ arr = rng.random((40, 12, 10)).astype("float32")
+ d = DynamicArray(zarr.array(arr, chunks=(8, 5, 4)))
+ out = str(tmp_path / "cs.zarr")
+ io.write(ops.cumsum(d, 0), out, zarr_format=2, chunks=(8, 5, 4), region_size_mb=0.002)
+ np.testing.assert_allclose(zarr.open(out, mode="r")[:], np.cumsum(arr, axis=0), atol=1e-3)
+
+
+def test_axis_out_of_range(data):
+ _, d = data
+ with pytest.raises(ValueError):
+ ops.cumsum(d, 3)
diff --git a/tests/test_transform_correctness.py b/tests/test_transform_correctness.py
new file mode 100644
index 0000000..ffd20d0
--- /dev/null
+++ b/tests/test_transform_correctness.py
@@ -0,0 +1,234 @@
+"""
+Correctness fuzz harness for the lazy Transform layer.
+
+The risky part of dyna_zarr is the pull-based ``read(key)`` slice-translation math
+inside every Transform: given an *output-space* slice, each transform must map it
+back to the correct *input-space* read and return exactly what NumPy would for the
+same operation followed by the same slice. This module fuzzes that contract.
+
+For every operation we check two things against a NumPy reference:
+
+1. **Full compute** ``op(da).compute()`` == ``op_np(arr)``
+2. **Random sub-slice** ``op(da)[k]`` for many random ``k`` == ``op_np(arr)[k]``
+
+(2) is where the ``read(key)`` composition math lives, so most ops are fuzzed there;
+ops whose ``read`` deliberately materializes the whole input (pad/tile/roll/flip/
+flatten) are still checked for correctness under sub-slicing even though they are not
+memory-bound -- their memory behaviour is a *performance* concern tracked in the
+non-shipped benchmark report, not a correctness bug.
+"""
+
+import numpy as np
+import pytest
+import zarr
+
+from dyna_zarr import DynamicArray, operations
+
+
+def _random_key_signed(rng, shape):
+ """Random basic-index key that INCLUDES negative starts/stops/ints (regression guard
+ for SliceTransform, which previously mis-sized negative slices like a[:-1])."""
+ key = []
+ for size in shape:
+ if size == 0:
+ key.append(slice(None))
+ continue
+ c = rng.integers(0, 5)
+ if c == 0:
+ i = int(rng.integers(-size, size))
+ key.append(i)
+ elif c == 1:
+ key.append(slice(None))
+ else:
+ lo = int(rng.integers(-size - 1, size + 1))
+ hi = int(rng.integers(-size - 1, size + 1))
+ step = int(rng.integers(1, 3))
+ key.append(slice(lo, hi, step))
+ return tuple(key)
+
+
+@pytest.mark.parametrize("chunks", [(4, 5, 6), (2, 2, 2), (1, 5, 6)])
+def test_slice_negative_indices(chunks):
+ """da[k] == arr[k] for keys with negative indices, single and composed (slice-of-slice)."""
+ rng = np.random.default_rng(2024)
+ arr = rng.random((4, 5, 6))
+ da = DynamicArray(zarr.array(arr, chunks=chunks))
+ for _ in range(200):
+ k = _random_key_signed(rng, arr.shape)
+ got = np.asarray(da[k].compute())
+ ref = arr[k]
+ assert got.shape == ref.shape, f"key {k}: {got.shape} != {ref.shape}"
+ np.testing.assert_array_equal(got, ref, err_msg=f"key {k}")
+ # compose a second (also-signed) slice on top
+ if got.ndim:
+ k2 = _random_key_signed(rng, got.shape)
+ np.testing.assert_array_equal(np.asarray(da[k][k2].compute()), ref[k2],
+ err_msg=f"composed {k} then {k2}")
+
+
+# --------------------------------------------------------------------------- #
+# Helpers
+# --------------------------------------------------------------------------- #
+
+def da_from_numpy(arr: np.ndarray) -> DynamicArray:
+ """Wrap a NumPy array as a DynamicArray via an in-memory zarr array."""
+ z = zarr.array(arr, chunks=arr.shape)
+ return DynamicArray(z)
+
+
+def random_key(rng: np.random.Generator, shape: tuple) -> tuple:
+ """Generate a random valid basic-indexing key for an array of ``shape``.
+
+ Mixes integers, full slices, bounded slices, and stepped slices per axis.
+ """
+ key = []
+ for size in shape:
+ choice = rng.integers(0, 5)
+ if choice == 0: # integer index
+ key.append(int(rng.integers(0, size)))
+ elif choice == 1: # full slice
+ key.append(slice(None))
+ elif choice == 2: # start:stop
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ key.append(slice(int(a), int(b)))
+ elif choice == 3: # start:stop:step
+ a, b = sorted(rng.integers(0, size + 1, size=2))
+ step = int(rng.integers(1, 3))
+ key.append(slice(int(a), int(b), step))
+ else: # open-ended
+ a = int(rng.integers(0, size + 1))
+ key.append(slice(a, None))
+ return tuple(key)
+
+
+def assert_close(got, ref, label):
+ got = np.asarray(got)
+ ref = np.asarray(ref)
+ assert got.shape == ref.shape, f"{label}: shape {got.shape} != {ref.shape}"
+ if np.issubdtype(ref.dtype, np.floating):
+ np.testing.assert_allclose(got, ref, rtol=1e-6, atol=1e-6,
+ err_msg=f"{label}: value mismatch")
+ else:
+ np.testing.assert_array_equal(got, ref, err_msg=f"{label}: value mismatch")
+
+
+# --------------------------------------------------------------------------- #
+# Operation registry: (name, dyna_fn, numpy_fn, supports_subslice)
+#
+# dyna_fn / numpy_fn take the *primary* operand (DynamicArray / ndarray) and,
+# for binary ops, close over a second operand built the same way.
+# --------------------------------------------------------------------------- #
+
+BASE_SHAPE = (4, 5, 6)
+
+
+def build_registry(arr, da, arr2, da2):
+ """Return list of op specs bound to concrete operands."""
+ perm = (2, 0, 1)
+ return [
+ # name, dyna_fn, numpy_fn, subslice
+ ("transpose", lambda: operations.transpose(da, perm), lambda: np.transpose(arr, perm), True),
+ ("swap_axes", lambda: operations.swap_axes(da, 0, 2), lambda: np.swapaxes(arr, 0, 2), True),
+ ("expand_dims", lambda: operations.expand_dims(da, 1), lambda: np.expand_dims(arr, 1), True),
+ ("squeeze", lambda: operations.squeeze(
+ operations.expand_dims(da, 1), 1), lambda: np.squeeze(np.expand_dims(arr, 1), 1), True),
+ ("flip", lambda: operations.flip(da, 1), lambda: np.flip(arr, 1), True),
+ ("roll", lambda: operations.roll(da, 2, 0), lambda: np.roll(arr, 2, 0), True),
+ ("pad", lambda: operations.pad(da, 1), lambda: np.pad(arr, 1), True),
+ ("tile", lambda: operations.tile(da, (1, 2, 1)), lambda: np.tile(arr, (1, 2, 1)), True),
+ ("concatenate", lambda: operations.concatenate([da, da2], 0), lambda: np.concatenate([arr, arr2], 0), True),
+ ("stack", lambda: operations.stack([da, da2], 0), lambda: np.stack([arr, arr2], 0), True),
+ ("clip", lambda: operations.clip(da, 0.2, 0.8), lambda: np.clip(arr, 0.2, 0.8), True),
+ ("abs", lambda: operations.abs(da), lambda: np.abs(arr), True),
+ ("sign", lambda: operations.sign(da), lambda: np.sign(arr), True),
+ ("round", lambda: operations.round(da, 2), lambda: np.round(arr, 2), True),
+ ("sqrt", lambda: operations.sqrt(da), lambda: np.sqrt(np.abs(arr)), True),
+ ("multiply", lambda: operations.multiply(da, da2), lambda: arr * arr2, True),
+ ("add", lambda: operations.add(da, da2), lambda: arr + arr2, True),
+ ("multiply_scalar", lambda: operations.multiply(da, 3.0), lambda: arr * 3.0, True),
+ ("where", lambda: operations.where(da, da, da2), lambda: np.where(arr, arr, arr2), True),
+ # --- map_blocks ufunc surface: unary (da2 is strictly positive, safe for log/recip) ---
+ ("negative", lambda: operations.negative(da), lambda: np.negative(arr), True),
+ ("square", lambda: operations.square(da), lambda: np.square(arr), True),
+ ("exp", lambda: operations.exp(da), lambda: np.exp(arr), True),
+ ("log", lambda: operations.log(da2), lambda: np.log(arr2), True),
+ ("log2", lambda: operations.log2(da2), lambda: np.log2(arr2), True),
+ ("log10", lambda: operations.log10(da2), lambda: np.log10(arr2), True),
+ ("floor", lambda: operations.floor(da2), lambda: np.floor(arr2), True),
+ ("ceil", lambda: operations.ceil(da2), lambda: np.ceil(arr2), True),
+ ("reciprocal", lambda: operations.reciprocal(da2), lambda: np.reciprocal(arr2), True),
+ ("astype_i16", lambda: operations.astype(da2, np.int16), lambda: arr2.astype(np.int16), True),
+ # --- binary ---
+ ("subtract", lambda: operations.subtract(da, da2), lambda: np.subtract(arr, arr2), True),
+ ("divide", lambda: operations.divide(da, da2), lambda: np.divide(arr, arr2), True),
+ ("floor_divide", lambda: operations.floor_divide(da, da2), lambda: np.floor_divide(arr, arr2), True),
+ ("mod", lambda: operations.mod(da, da2), lambda: np.mod(arr, arr2), True),
+ ("power", lambda: operations.power(da2, da2), lambda: np.power(arr2, arr2), True),
+ ("maximum", lambda: operations.maximum(da, da2), lambda: np.maximum(arr, arr2), True),
+ ("minimum", lambda: operations.minimum(da, da2), lambda: np.minimum(arr, arr2), True),
+ ("add_scalar", lambda: operations.add(da, 2.5), lambda: arr + 2.5, True),
+ # --- comparisons & logical (bool out) ---
+ ("greater", lambda: operations.greater(da, da2), lambda: np.greater(arr, arr2), True),
+ ("less_equal", lambda: operations.less_equal(da, da2), lambda: np.less_equal(arr, arr2), True),
+ ("equal", lambda: operations.equal(da, da), lambda: np.equal(arr, arr), True),
+ ("logical_and", lambda: operations.logical_and(da, da2), lambda: np.logical_and(arr, arr2), True),
+ ("logical_not", lambda: operations.logical_not(da), lambda: np.logical_not(arr), True),
+ # reductions: full-compute only (sub-slicing a reduced axis is a separate contract)
+ ("min_axis0", lambda: operations.min(da, 0), lambda: np.min(arr, 0), False),
+ ("max_axis0", lambda: operations.max(da, 0), lambda: np.max(arr, 0), False),
+ ]
+
+
+@pytest.fixture
+def operands():
+ rng = np.random.default_rng(0)
+ arr = rng.random(BASE_SHAPE) # float64 in [0, 1)
+ arr[arr < 0.05] = 0.0 # a few exact zeros for sign/where
+ arr2 = rng.random(BASE_SHAPE) + 0.1
+ return arr, da_from_numpy(arr), arr2, da_from_numpy(arr2)
+
+
+def op_ids(reg):
+ return [spec[0] for spec in reg]
+
+
+# --------------------------------------------------------------------------- #
+# Tests
+# --------------------------------------------------------------------------- #
+
+def _registry(operands):
+ return build_registry(*operands)
+
+
+def test_full_compute_matches_numpy(operands):
+ """op(da).compute() == op_np(arr) for every operation."""
+ reg = _registry(operands)
+ failures = []
+ for name, dyna_fn, np_fn, _ in reg:
+ try:
+ assert_close(dyna_fn().compute(), np_fn(), f"{name}[full]")
+ except Exception as e:
+ failures.append(f"{name}[full]: {type(e).__name__}: {str(e).splitlines()[0]}")
+ assert not failures, "Full-compute mismatches:\n" + "\n".join(failures)
+
+
+@pytest.mark.parametrize("seed", [0, 1, 7, 1234, 99999])
+def test_random_subslice_matches_numpy(operands, seed):
+ """op(da)[k] == op_np(arr)[k] over many random keys (the read(key) math)."""
+ rng = np.random.default_rng(seed)
+ reg = _registry(operands)
+ failures = []
+ for name, dyna_fn, np_fn, subslice in reg:
+ if not subslice:
+ continue
+ ref_full = np_fn()
+ result = dyna_fn()
+ for _ in range(40):
+ k = random_key(rng, result.shape)
+ try:
+ got = result[k].compute()
+ assert_close(got, ref_full[k], f"{name}[{k}]")
+ except Exception as e:
+ failures.append(f"{name}[{k}]: {type(e).__name__}: {str(e).splitlines()[0]}")
+ break # one failure per op is enough signal
+ assert not failures, "Sub-slice mismatches:\n" + "\n".join(failures)