diff --git a/README.md b/README.md index 5af451d..ccbdb86 100644 --- a/README.md +++ b/README.md @@ -45,10 +45,12 @@ from infinite_tensor import InfiniteTensor, TensorWindow, MemoryTileStore tile_store = MemoryTileStore() + def your_processing_function(ctx): # ctx is the window index (e.g. (wy, wx) for 2D), NOT pixel coordinates. return torch.ones(512, 512) + tensor = InfiniteTensor( shape=(None, None), f=your_processing_function, @@ -72,6 +74,28 @@ result = tensor[0:1024, 0:1024] `__getitem__` supports integer indices, slices, and negative indices. Integer indices squeeze the dimension; slices preserve it. +To read several regions at once, build each region with `sl` and index with a list of them. You get back a list of tensors in the same order: + +```python +from infinite_tensor import sl + +a, b = tensor[[sl[0:512, 0:512], sl[4096:4608, 0:512]]] +``` + +`sl[...]` creates a region for batched reads using normal indexing syntax: `sl[0:512, 5]` is the region `tensor[0:512, 5]` would read. + +All windows needed by every region are computed in one pass (deduplicated, and batched together when `batch_size` is set) and stay cached until every region has been read. Any list passed to `[]` is treated as a list of regions. + +To get results as they are computed instead of all at once, use `iter_regions`. It processes windows in region order in full batches of the tensor's `batch_size`, and after each batch yields a list of the regions that are now complete: + +```python +for chunk in tensor.iter_regions(sl[i : i + 512, 0:512] for i in range(0, 8192, 512)): + for result in chunk: + ... +``` + +Results come in the same order as the regions, so concatenating the chunks gives `tensor[regions]`. With no `batch_size` you get a single chunk; with `batch_size=1` you get roughly one region at a time. `regions` can be a lazy generator. + ### 3. Device and dtype Pass `device=` and/or `dtype=` at construction. `f` must return tensors matching both exactly (`DeviceMismatchError` / `DtypeMismatchError` otherwise). Defaults are `cpu` / `torch.float32`. Upstream arg slices are **not** auto-transferred; if an upstream lives elsewhere, move/cast it inside `f`. @@ -85,7 +109,7 @@ Pass `device=` and/or `dtype=` at construction. `f` must return tensors matching ```python tile_store = MemoryTileStore( cache_size_bytes=200 * 1024 * 1024, # 200 MB - cache_size_windows=None, # no window-count limit + cache_size_windows=None, # no window-count limit ) ``` @@ -124,9 +148,11 @@ WINDOW = 512 PADDING = 2 PADDED = WINDOW + 2 * PADDING + def random_image(ctx): return torch.rand(3, WINDOW, WINDOW) + base = InfiniteTensor( shape=(3, None, None), f=random_image, @@ -141,10 +167,12 @@ padded_input_window = TensorWindow( offset=(0, -PADDING, -PADDING), ) + def blur(ctx, padded_input): blurred = F.gaussian_blur(padded_input, kernel_size=5, sigma=1.0) return blurred[:, PADDING:-PADDING, PADDING:-PADDING] + blurred = InfiniteTensor( shape=(3, None, None), f=blur, @@ -170,11 +198,12 @@ Setting `batch_size=N` changes the contract of `f`: ```python def blur_batched(ctxs, padded_inputs): - stack = torch.stack(padded_inputs) # (B, 3, PADDED, PADDED) + stack = torch.stack(padded_inputs) # (B, 3, PADDED, PADDED) blurred = F.gaussian_blur(stack, kernel_size=5, sigma=1.0) cropped = blurred[:, :, PADDING:-PADDING, PADDING:-PADDING] return [cropped[i] for i in range(cropped.shape[0])] + blurred = InfiniteTensor( shape=(3, None, None), f=blur_batched, @@ -220,10 +249,12 @@ Overlapping windows are summed by default. Two optional `InfiniteTensor` kwargs import torch from infinite_tensor import InfiniteTensor, TensorWindow, MemoryTileStore + def f(ctx): wy, wx = ctx return torch.full((64, 64), float(wy + wx)) + tensor = InfiniteTensor( shape=(None, None), f=f, @@ -274,7 +305,7 @@ HDF5TileStore( "tensor_data.h5", tile_size=512, cache_size_bytes=256 * 1024 * 1024, # 256 MiB cache; default 100 MiB - cache_size_tiles=None, # no tile-count limit + cache_size_tiles=None, # no tile-count limit ) ``` @@ -284,12 +315,12 @@ Reconnecting in a later process: with HDF5TileStore("tensor_data.h5", tile_size=512) as tile_store: tensor = InfiniteTensor( shape=(None, None), - f=your_processing_function, # same f, or one that produces the same results + f=your_processing_function, # same f, or one that produces the same results output_window=TensorWindow((512, 512)), tile_store=tile_store, tensor_id="my_infinite_tensor", # same id as before ) - result = tensor[0:2048, 0:2048] # served from disk; f is only called for new windows + result = tensor[0:2048, 0:2048] # served from disk; f is only called for new windows ``` Once data is saved to a persistent store, `clear_cache` only drops in-memory buffers. Use `tile_store.clear_tensor("my_infinite_tensor")` to fully wipe the saved state for that tensor. diff --git a/infinite_tensor/__init__.py b/infinite_tensor/__init__.py index 7f3945f..44f962d 100644 --- a/infinite_tensor/__init__.py +++ b/infinite_tensor/__init__.py @@ -11,6 +11,7 @@ ) from .tensor_window import TensorWindow from .tilestore import MemoryTileStore +from .utils import sl # Optional HDF5 support try: @@ -35,6 +36,7 @@ "ValidationError", "TensorWindow", "MemoryTileStore", + "sl", ] if _HAS_HDF5: diff --git a/infinite_tensor/infinite_tensor.py b/infinite_tensor/infinite_tensor.py index 01e1ee1..c75a6c8 100644 --- a/infinite_tensor/infinite_tensor.py +++ b/infinite_tensor/infinite_tensor.py @@ -3,7 +3,8 @@ import itertools import logging import uuid -from collections.abc import Iterable +from collections import deque +from collections.abc import Iterable, Iterator from typing import Any, Callable import torch @@ -403,23 +404,132 @@ def to(self, *args, **kwargs) -> InfiniteTensor: raise return self - def __getitem__(self, indices: tuple[int | slice, ...]) -> torch.Tensor: - """Return tensor values for the requested pixel-space indices.""" - pixel_slices, collapse_dims = standardize_indices(self.shape, indices) - pixel_slices_tuple: tuple[slice, ...] = tuple(pixel_slices) - logger.debug(f"Accessing tensor slice with indices: {pixel_slices_tuple}") + def __getitem__(self, indices: Any) -> torch.Tensor | list[torch.Tensor]: + """Return tensor values for the requested pixel-space indices. + + ``indices`` is either a single region (ints / slices / ellipsis, as + with a normal tensor) or a ``list`` of regions, e.g. + ``tensor[[sl[0:512, 0:512], sl[4096:4608, 0:512]]]``. A list always + means a list of regions and returns a ``list[torch.Tensor]`` in the + same order; every window needed by any region is computed in one pass + (deduplicated and batched together) and kept cached until all regions + have been read. + + Use ``sl[...]`` (``from infinite_tensor import sl``) to create regions + for batched reads: ``sl[0:512, 5]`` is the region ``tensor[0:512, 5]`` + would read. + """ + is_region_list = isinstance(indices, list) + regions = indices if is_region_list else [indices] + + standardized = [self._standardize_region(region) for region in regions] + pixel_regions = [pixel_region for pixel_region, _ in standardized] + logger.debug(f"Accessing tensor slices with indices: {pixel_regions}") + + self._store.begin_access(self._uuid) + try: + self._ensure_processed_range(pixel_regions) + results = self._read_processed_regions(standardized) + finally: + self._store.end_access(self._uuid) + return results if is_region_list else results[0] + + def iter_regions(self, regions: Iterable[Any]) -> Iterator[list[torch.Tensor]]: + """Read ``regions`` incrementally, yielding results as their windows are computed. + + Windows are processed in region order, in full batches of + ``batch_size``. After each batch, every leading region whose windows + are all processed is read and yielded as a list, so concatenating the + yielded lists gives the same result as ``tensor[list(regions)]``. + With ``batch_size=None`` everything is processed in one pass and + yielded as a single list; with ``batch_size=1`` results arrive about + one region at a time. + + ``regions`` may be any iterable (including a lazy generator) of + regions, e.g. ``tensor.iter_regions(sl[0, i:i + 64] for i in range(0, 1024, 64))``. + Regions that are already cached are yielded with the next batch. + """ + waiting: deque[tuple[tuple[tuple[slice, ...], list[int]], int]] = deque() + pending: list[tuple[int, ...]] = [] + pending_set: set[tuple[int, ...]] = set() + batch_id = 0 # id of the batch ``pending`` is filling; lower ids are processed + + def run_batch_if_full() -> Iterator[list[torch.Tensor]]: + nonlocal pending, pending_set, batch_id + if len(pending) == self.batch_size: + results = self._run_batch_and_read_completed(pending, waiting, batch_id + 1) + pending, pending_set = [], set() + batch_id += 1 + if results: + yield results + + for region in regions: + standardized = self._standardize_region(region) + needs_batch = -1 + for window_index in self.output_window.intersecting_windows( + standardized[0], tensor_shape=self.shape + ): + if window_index in pending_set: + needs_batch = batch_id + elif not self._store.is_window_processed(self._uuid, window_index): + yield from run_batch_if_full() + pending.append(window_index) + pending_set.add(window_index) + needs_batch = batch_id + waiting.append((standardized, needs_batch)) + # Check again here so a region whose last window fills the batch is yielded right away. + yield from run_batch_if_full() + + if pending or waiting: + results = self._run_batch_and_read_completed(pending, waiting, batch_id + 1) + if results: + yield results + + def _run_batch_and_read_completed( + self, + pending: list[tuple[int, ...]], + waiting: deque[tuple[tuple[tuple[slice, ...], list[int]], int]], + processed_batches: int, + ) -> list[torch.Tensor]: + """Process ``pending``, then pop and read the leading ``waiting`` regions that are complete.""" + ready = [] + while waiting and waiting[0][1] < processed_batches: + ready.append(waiting.popleft()[0]) self._store.begin_access(self._uuid) try: - self._ensure_processed_range([pixel_slices_tuple]) - result = self._store.read_pixels(self._uuid, pixel_slices_tuple) + self._ensure_processed(pending) + # Re-check ready regions: windows from earlier batches may have been evicted. + self._ensure_processed_range([pixel_region for pixel_region, _ in ready]) + return self._read_processed_regions(ready) finally: self._store.end_access(self._uuid) - if collapse_dims: - target_shape = tuple(s for i, s in enumerate(result.shape) if i not in collapse_dims) - result = result.reshape(target_shape) - return result + def _standardize_region(self, region: Any) -> tuple[tuple[slice, ...], list[int]]: + """Normalize one user region to ``(pixel_slices, collapse_dims)``.""" + pixel_slices, collapse_dims = standardize_indices(self.shape, region) + return tuple(pixel_slices), collapse_dims + + def _read_processed_regions( + self, regions: list[tuple[tuple[slice, ...], list[int]]] + ) -> list[torch.Tensor]: + """Read standardized regions whose windows are all already processed. + + Does not compute anything: the caller must have run + ``_ensure_processed_range`` on these regions within the same + ``begin_access`` / ``end_access`` bracket. Integer-indexed dims are + squeezed from each result. + """ + results = [] + for pixel_region, collapse_dims in regions: + result = self._store.read_pixels(self._uuid, pixel_region) + if collapse_dims: + target_shape = tuple( + s for d, s in enumerate(result.shape) if d not in collapse_dims + ) + result = result.reshape(target_shape) + results.append(result) + return results def _ensure_processed_range(self, pixel_ranges: Iterable[tuple[slice, ...]]) -> None: """Ensure every window intersecting ``pixel_ranges`` is processed in the store.""" diff --git a/infinite_tensor/utils.py b/infinite_tensor/utils.py index 15c401a..bdc1772 100644 --- a/infinite_tensor/utils.py +++ b/infinite_tensor/utils.py @@ -7,6 +7,8 @@ from __future__ import annotations +from typing import Any + def normalize_slice(idx: int | slice, shape_dim: int | None): """Convert an index or slice to a normalized slice format. @@ -81,10 +83,14 @@ def standardize_indices( Raises: IndexError: If indices are invalid for the tensor shape """ - if isinstance(indices, list): - indices = tuple(indices) if not isinstance(indices, tuple): indices = (indices,) + for idx in indices: + if not isinstance(idx, (int, slice)) and idx is not Ellipsis: + raise TypeError( + f"Invalid index {idx!r}: each index must be an int, slice, or Ellipsis " + "(to read several regions, pass a list of regions, e.g. [sl[0:8], sl[16:24]])" + ) # Handle ellipsis by expanding it to the appropriate number of full slices if Ellipsis in indices: @@ -148,3 +154,19 @@ def standardize_indices( ] return normalized_slices, collapse_dims + + +class _RegionBuilder: + """Build a region from bracket syntax: ``sl[0:512, 5]`` -> ``(slice(0, 512, None), 5)``. + + Always returns a tuple, so ``sl[0:8]`` is a valid one-dimensional region. + """ + + def __getitem__(self, indices: Any) -> tuple: + return indices if isinstance(indices, tuple) else (indices,) + + def __repr__(self) -> str: + return "sl" + + +sl = _RegionBuilder() diff --git a/pyproject.toml b/pyproject.toml index ff92be1..be5c87f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "infinite-tensor" -version = "0.3.0" +version = "0.4.0" description = "A Python library for performing operations on tensors with infinite dimensions." readme = "README.md" requires-python = ">=3.9" diff --git a/tests/test_infinite_tensor.py b/tests/test_infinite_tensor.py index 80705dc..46fc9a5 100644 --- a/tests/test_infinite_tensor.py +++ b/tests/test_infinite_tensor.py @@ -1,5 +1,6 @@ """Comprehensive tests for InfiniteTensor functionality using pytest.""" +import itertools import uuid import pytest @@ -10,6 +11,7 @@ InfiniteTensor, TensorWindow, ValidationError, + sl, ) from infinite_tensor.tilestore import MemoryTileStore @@ -646,3 +648,171 @@ def expected_at(row, col): expected = torch.tensor([[float(expected_at(r, c)) for c in range(8)] for r in range(8)]) assert torch.equal(result, expected) + + +class TestRegionListIndexing: + """Indexing with a list of regions returns one tensor per region.""" + + @staticmethod + def _coord_tensor(store, batch_size=None, calls=None): + def f(ctx): + if calls is not None: + calls.append(ctx) + if batch_size is not None: + return [torch.full((1, 32, 32), float(c[1] * 100 + c[2])) for c in ctx] + return torch.full((1, 32, 32), float(ctx[1] * 100 + ctx[2])) + + return InfiniteTensor( + (1, None, None), + f, + TensorWindow((1, 32, 32)), + tile_store=store, + batch_size=batch_size, + ) + + def test_region_list_matches_individual_reads(self, tile_store): + tensor = self._coord_tensor(tile_store) + regions = [ + sl[:, 0:40, 10:70], + sl[0, -32:5, 100:132], + sl[0, 33, 0:64:3], + ] + results = tensor[regions] + + assert isinstance(results, list) + assert len(results) == len(regions) + for region, result in zip(regions, results): + assert torch.equal(result, tensor[region]) + assert results[0].shape == (1, 40, 60) + assert results[1].shape == (37, 32) + assert results[2].shape == (22,) + + def test_region_list_computes_shared_windows_once(self, tile_store): + calls = [] + tensor = self._coord_tensor(tile_store, calls=calls) + tensor[[sl[0, 0:64, 0:32], sl[0, 32:96, 0:32]]] + assert sorted(calls) == [(0, 0, 0), (0, 1, 0), (0, 2, 0)] + + def test_region_list_batches_across_regions(self, tile_store): + calls = [] + tensor = self._coord_tensor(tile_store, batch_size=8, calls=calls) + tensor[[sl[0, 0:32, 0:32], sl[0, 320:352, 320:352]]] + assert len(calls) == 1 + assert sorted(calls[0]) == [(0, 0, 0), (0, 10, 10)] + + def test_region_list_held_in_cache_until_all_read(self): + store = MemoryTileStore(cache_size_windows=1) + calls = [] + tensor = self._coord_tensor(store, calls=calls) + a, b = tensor[[sl[0, 0:32, 0:32], sl[0, 32:64, 0:32]]] + assert torch.all(a == 0) and torch.all(b == 100) + assert len(calls) == 2 + assert len(store._windows) == 1 + + def test_empty_region_list(self, tile_store): + tensor = self._coord_tensor(tile_store) + assert tensor[[]] == [] + + def test_any_list_is_region_list(self, tile_store): + """Bare ints / slices in a list are each their own region, not one combined index.""" + tensor = InfiniteTensor( + (None,), + lambda ctx: torch.arange(8, dtype=torch.float32) + 8 * ctx[0], + TensorWindow((8,)), + tile_store=tile_store, + ) + a, b = tensor[[3, slice(10, 14)]] + assert a.shape == () and a.item() == 3.0 + assert torch.equal(b, torch.arange(10, 14, dtype=torch.float32)) + + def test_one_dimensional_regions(self, tile_store): + tensor = InfiniteTensor( + (None,), + lambda ctx: torch.full((8,), float(ctx[0])), + TensorWindow((8,)), + tile_store=tile_store, + ) + a, b, c = tensor[[sl[0:4], sl[16:24], 9]] + assert torch.equal(a, torch.zeros(4)) + assert torch.equal(b, torch.full((8,), 2.0)) + assert c.shape == () and c.item() == 1.0 + + def test_sl_always_returns_tuple(self): + assert sl[0:8] == (slice(0, 8),) + assert sl[3] == (3,) + assert sl[0:8, 5] == (slice(0, 8), 5) + assert sl[..., 0] == (Ellipsis, 0) + + def test_nested_list_raises_type_error(self, tile_store): + tensor = self._coord_tensor(tile_store) + with pytest.raises(TypeError, match="Invalid index"): + tensor[[[0, slice(0, 8), slice(0, 8)]]] + + +class TestIterRegions: + """iter_regions yields region results incrementally, following batch_size.""" + + @staticmethod + def _row_tensor(store, batch_size=None, calls=None): + """1D tensor of 8-wide windows; window i is filled with i.""" + + def f(ctx): + if calls is not None: + calls.append(ctx) + if batch_size is not None: + return [torch.full((8,), float(c[0])) for c in ctx] + return torch.full((8,), float(ctx[0])) + + return InfiniteTensor( + (None,), f, TensorWindow((8,)), tile_store=store, batch_size=batch_size + ) + + REGIONS = [sl[0:8], sl[8:24], sl[20:28], sl[3], sl[40:48]] + + @pytest.mark.parametrize("batch_size", [None, 1, 2, 3, 100]) + def test_concatenated_yields_match_getitem(self, tile_store, batch_size): + tensor = self._row_tensor(tile_store, batch_size=batch_size) + streamed = [r for chunk in tensor.iter_regions(self.REGIONS) for r in chunk] + expected = self._row_tensor(MemoryTileStore())[self.REGIONS] + assert len(streamed) == len(expected) + for a, b in zip(streamed, expected): + assert torch.equal(a, b) + + def test_no_batch_size_yields_single_list(self, tile_store): + tensor = self._row_tensor(tile_store) + chunks = list(tensor.iter_regions(self.REGIONS)) + assert [len(c) for c in chunks] == [len(self.REGIONS)] + + def test_batch_size_one_yields_region_by_region(self, tile_store): + tensor = self._row_tensor(tile_store, batch_size=1) + # Windows per region: [0], [1, 2], [2, 3], [0], [5] + # sl[3] only needs window 0 (already done), so it rides along with the next batch. + chunks = list(tensor.iter_regions(self.REGIONS)) + assert [len(c) for c in chunks] == [1, 1, 1, 2] + + def test_batches_stay_full_across_regions(self, tile_store): + calls = [] + tensor = self._row_tensor(tile_store, batch_size=2, calls=calls) + chunks = list(tensor.iter_regions(self.REGIONS)) + assert calls == [[(0,), (1,)], [(2,), (3,)], [(5,)]] + assert [len(c) for c in chunks] == [1, 2, 2] + + def test_is_lazy(self, tile_store): + calls = [] + tensor = self._row_tensor(tile_store, batch_size=1, calls=calls) + it = tensor.iter_regions(sl[i * 8 : i * 8 + 8] for i in itertools.count()) + first = next(it) + assert torch.equal(first[0], torch.zeros(8)) + assert calls == [[(0,)]] + + def test_recomputes_evicted_windows(self): + store = MemoryTileStore(cache_size_windows=1) + tensor = self._row_tensor(store, batch_size=1) + # sl[0:16] spans windows 0 and 1, processed in separate batches; window 0 + # is evicted after the first batch and must be recomputed before reading. + (result,) = list(tensor.iter_regions([sl[0:16]]))[0] + assert torch.equal(result, torch.cat([torch.zeros(8), torch.ones(8)])) + + def test_empty(self, tile_store): + tensor = self._row_tensor(tile_store) + assert list(tensor.iter_regions([])) == []