Region-list indexing and incremental iter_regions

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.

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:

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."""

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()

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"