bioimageio.core.tensor

  1from __future__ import annotations
  2
  3import collections.abc
  4from itertools import permutations
  5from typing import (
  6    TYPE_CHECKING,
  7    Any,
  8    Callable,
  9    Dict,
 10    Iterator,
 11    Mapping,
 12    Optional,
 13    Sequence,
 14    Tuple,
 15    Union,
 16    cast,
 17    get_args,
 18)
 19
 20import numpy as np
 21import xarray as xr
 22from loguru import logger
 23from numpy.typing import DTypeLike, NDArray
 24from typing_extensions import Self, assert_never
 25
 26from bioimageio.spec.model import v0_5
 27
 28from ._magic_tensor_ops import MagicTensorOpsMixin
 29from .axis import Axis, AxisId, AxisInfo, AxisLike, PerAxis
 30from .common import (
 31    CropWhere,
 32    DTypeStr,
 33    PadMode,
 34    PadWhere,
 35    PadWidth,
 36    PadWidthLike,
 37    SliceInfo,
 38)
 39
 40if TYPE_CHECKING:
 41    from numpy.typing import ArrayLike, NDArray
 42
 43
 44_ScalarOrArray = Union["ArrayLike", np.generic, "NDArray[Any]"]  # TODO: add "DaskArray"
 45
 46
 47# TODO: complete docstrings
 48class Tensor(MagicTensorOpsMixin):
 49    """A wrapper around an xr.DataArray for better integration with bioimageio.spec
 50    and improved type annotations."""
 51
 52    _Compatible = Union["Tensor", xr.DataArray, _ScalarOrArray]
 53
 54    def __init__(
 55        self,
 56        array: NDArray[Any],
 57        dims: Sequence[Union[AxisId, AxisLike]],
 58    ) -> None:
 59        super().__init__()
 60        axes = tuple(
 61            a if isinstance(a, AxisId) else AxisInfo.create(a).id for a in dims
 62        )
 63        self._data = xr.DataArray(array, dims=axes)
 64
 65    def __array__(self, dtype: DTypeLike = None):
 66        return np.asarray(self._data, dtype=dtype)
 67
 68    def __getitem__(
 69        self, key: Union[SliceInfo, slice, int, PerAxis[Union[SliceInfo, slice, int]]]
 70    ) -> Self:
 71        if isinstance(key, SliceInfo):
 72            key = slice(*key)
 73        elif isinstance(key, collections.abc.Mapping):
 74            key = {
 75                a: s if isinstance(s, int) else s if isinstance(s, slice) else slice(*s)
 76                for a, s in key.items()
 77            }
 78        return self.__class__.from_xarray(self._data[key])
 79
 80    def __setitem__(self, key: PerAxis[Union[SliceInfo, slice]], value: Tensor) -> None:
 81        key = {a: s if isinstance(s, slice) else slice(*s) for a, s in key.items()}
 82        self._data[key] = value._data
 83
 84    def __len__(self) -> int:
 85        return len(self.data)
 86
 87    def _iter(self: Any) -> Iterator[Any]:
 88        for n in range(len(self)):
 89            yield self[n]
 90
 91    def __iter__(self: Any) -> Iterator[Any]:
 92        if self.ndim == 0:
 93            raise TypeError("iteration over a 0-d array")
 94        return self._iter()
 95
 96    def _binary_op(
 97        self,
 98        other: _Compatible,
 99        f: Callable[[Any, Any], Any],
100        reflexive: bool = False,
101    ) -> Self:
102        data = self._data._binary_op(  # pyright: ignore[reportPrivateUsage]
103            (other._data if isinstance(other, Tensor) else other),
104            f,
105            reflexive,
106        )
107        return self.__class__.from_xarray(data)
108
109    def _inplace_binary_op(
110        self,
111        other: _Compatible,
112        f: Callable[[Any, Any], Any],
113    ) -> Self:
114        _ = self._data._inplace_binary_op(  # pyright: ignore[reportPrivateUsage]
115            (
116                other_d
117                if (other_d := getattr(other, "data")) is not None
118                and isinstance(
119                    other_d,
120                    xr.DataArray,
121                )
122                else other
123            ),
124            f,
125        )
126        return self
127
128    def _unary_op(self, f: Callable[[Any], Any], *args: Any, **kwargs: Any) -> Self:
129        data = self._data._unary_op(  # pyright: ignore[reportPrivateUsage]
130            f, *args, **kwargs
131        )
132        return self.__class__.from_xarray(data)
133
134    @classmethod
135    def from_xarray(cls, data_array: xr.DataArray) -> Self:
136        """create a `Tensor` from an xarray data array
137
138        note for internal use: this factory method is round-trip save
139            for any `Tensor`'s  `data` property (an xarray.DataArray).
140        """
141        return cls(
142            array=data_array.data, dims=tuple(AxisId(d) for d in data_array.dims)
143        )
144
145    @classmethod
146    def from_numpy(
147        cls,
148        array: NDArray[Any],
149        *,
150        dims: Optional[Union[AxisLike, Sequence[AxisLike]]],
151    ) -> Tensor:
152        """create a `Tensor` from a numpy array
153
154        Args:
155            array: the nd numpy array
156            axes: A description of the array's axes,
157                if None axes are guessed (which might fail and raise a ValueError.)
158
159        Raises:
160            ValueError: if `axes` is None and axes guessing fails.
161        """
162
163        if dims is None:
164            return cls._interprete_array_wo_known_axes(array)
165        elif isinstance(dims, (str, Axis, v0_5.AxisBase)):
166            dims = [dims]
167
168        axis_infos = [AxisInfo.create(a) for a in dims]
169        original_shape = tuple(array.shape)
170
171        successful_view = _get_array_view(array, axis_infos)
172        if successful_view is None:
173            raise ValueError(
174                f"Array shape {original_shape} does not map to axes {dims}"
175            )
176
177        return Tensor(successful_view, dims=tuple(a.id for a in axis_infos))
178
179    @property
180    def data(self):
181        return self._data
182
183    @property
184    def dims(self):  # TODO: rename to `axes`?
185        """Tuple of dimension names associated with this tensor."""
186        return cast(Tuple[AxisId, ...], self._data.dims)
187
188    @property
189    def tagged_shape(self):
190        """(alias for `sizes`) Ordered, immutable mapping from axis ids to lengths."""
191        return self.sizes
192
193    @property
194    def shape_tuple(self):
195        """Tuple of tensor axes lengths"""
196        return self._data.shape
197
198    @property
199    def size(self):
200        """Number of elements in the tensor.
201
202        Equal to math.prod(tensor.shape), i.e., the product of the tensors’ dimensions.
203        """
204        return self._data.size
205
206    def sum(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
207        """Reduce this Tensor's data by applying sum along some dimension(s)."""
208        return self.__class__.from_xarray(self._data.sum(dim=dim))
209
210    @property
211    def ndim(self):
212        """Number of tensor dimensions."""
213        return self._data.ndim
214
215    @property
216    def dtype(self) -> DTypeStr:
217        dt = str(self.data.dtype)  # pyright: ignore[reportUnknownArgumentType]
218        assert dt in get_args(DTypeStr)
219        return dt  # pyright: ignore[reportReturnType]
220
221    @property
222    def sizes(self):
223        """Ordered, immutable mapping from axis ids to axis lengths."""
224        return cast(Mapping[AxisId, int], self.data.sizes)
225
226    def astype(self, dtype: DTypeStr, *, copy: bool = False):
227        """Return tensor cast to `dtype`
228
229        note: if dtype is already satisfied copy if `copy`"""
230        return self.__class__.from_xarray(self._data.astype(dtype, copy=copy))
231
232    def clip(self, min: Optional[float] = None, max: Optional[float] = None):
233        """Return a tensor whose values are limited to [min, max].
234        At least one of max or min must be given."""
235        return self.__class__.from_xarray(self._data.clip(min, max))
236
237    def crop_to(
238        self,
239        sizes: PerAxis[int],
240        crop_where: Union[
241            CropWhere,
242            PerAxis[CropWhere],
243        ] = "left_and_right",
244    ) -> Self:
245        """crop to match `sizes`"""
246        if isinstance(crop_where, str):
247            crop_axis_where: PerAxis[CropWhere] = {a: crop_where for a in self.dims}
248        else:
249            crop_axis_where = crop_where
250
251        slices: Dict[AxisId, SliceInfo] = {}
252
253        for a, s_is in self.sizes.items():
254            if a not in sizes or sizes[a] == s_is:
255                pass
256            elif sizes[a] > s_is:
257                logger.warning(
258                    "Cannot crop axis {} of size {} to larger size {}",
259                    a,
260                    s_is,
261                    sizes[a],
262                )
263            elif a not in crop_axis_where:
264                raise ValueError(
265                    f"Don't know where to crop axis {a}, `crop_where`={crop_where}"
266                )
267            else:
268                crop_this_axis_where = crop_axis_where[a]
269                if crop_this_axis_where == "left":
270                    slices[a] = SliceInfo(s_is - sizes[a], s_is)
271                elif crop_this_axis_where == "right":
272                    slices[a] = SliceInfo(0, sizes[a])
273                elif crop_this_axis_where == "left_and_right":
274                    slices[a] = SliceInfo(
275                        start := (s_is - sizes[a]) // 2, sizes[a] + start
276                    )
277                else:
278                    assert_never(crop_this_axis_where)
279
280        return self[slices]
281
282    def expand_dims(self, dims: Union[Sequence[AxisId], PerAxis[int]]) -> Self:
283        return self.__class__.from_xarray(self._data.expand_dims(dims=dims))
284
285    def mean(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
286        return self.__class__.from_xarray(self._data.mean(dim=dim))
287
288    def std(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
289        return self.__class__.from_xarray(self._data.std(dim=dim))
290
291    def var(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
292        return self.__class__.from_xarray(self._data.var(dim=dim))
293
294    def pad(
295        self,
296        pad_width: PerAxis[PadWidthLike],
297        mode: PadMode = "symmetric",
298    ) -> Self:
299        pad_width = {a: PadWidth.create(p) for a, p in pad_width.items()}
300        return self.__class__.from_xarray(
301            self._data.pad(pad_width=pad_width, mode=mode)
302        )
303
304    def pad_to(
305        self,
306        sizes: PerAxis[int],
307        pad_where: Union[PadWhere, PerAxis[PadWhere]] = "left_and_right",
308        mode: PadMode = "symmetric",
309    ) -> Self:
310        """pad `tensor` to match `sizes`"""
311        if isinstance(pad_where, str):
312            pad_axis_where: PerAxis[PadWhere] = {a: pad_where for a in self.dims}
313        else:
314            pad_axis_where = pad_where
315
316        pad_width: Dict[AxisId, PadWidth] = {}
317        for a, s_is in self.sizes.items():
318            if a not in sizes or sizes[a] == s_is:
319                pad_width[a] = PadWidth(0, 0)
320            elif s_is > sizes[a]:
321                pad_width[a] = PadWidth(0, 0)
322                logger.warning(
323                    "Cannot pad axis {} of size {} to smaller size {}",
324                    a,
325                    s_is,
326                    sizes[a],
327                )
328            elif a not in pad_axis_where:
329                raise ValueError(
330                    f"Don't know where to pad axis {a}, `pad_where`={pad_where}"
331                )
332            else:
333                pad_this_axis_where = pad_axis_where[a]
334                d = sizes[a] - s_is
335                if pad_this_axis_where == "left":
336                    pad_width[a] = PadWidth(d, 0)
337                elif pad_this_axis_where == "right":
338                    pad_width[a] = PadWidth(0, d)
339                elif pad_this_axis_where == "left_and_right":
340                    pad_width[a] = PadWidth(left := d // 2, d - left)
341                else:
342                    assert_never(pad_this_axis_where)
343
344        return self.pad(pad_width, mode)
345
346    def quantile(
347        self,
348        q: Union[float, Sequence[float]],
349        dim: Optional[Union[AxisId, Sequence[AxisId]]] = None,
350    ) -> Self:
351        assert (
352            isinstance(q, (float, int))
353            and q >= 0.0
354            or not isinstance(q, (float, int))
355            and all(qq >= 0.0 for qq in q)
356        )
357        assert (
358            isinstance(q, (float, int))
359            and q <= 1.0
360            or not isinstance(q, (float, int))
361            and all(qq <= 1.0 for qq in q)
362        )
363        assert dim is None or (
364            (quantile_dim := AxisId("quantile")) != dim and quantile_dim not in set(dim)
365        )
366        return self.__class__.from_xarray(self._data.quantile(q, dim=dim))
367
368    def resize_to(
369        self,
370        sizes: PerAxis[int],
371        *,
372        pad_where: Union[
373            PadWhere,
374            PerAxis[PadWhere],
375        ] = "left_and_right",
376        crop_where: Union[
377            CropWhere,
378            PerAxis[CropWhere],
379        ] = "left_and_right",
380        pad_mode: PadMode = "symmetric",
381    ):
382        """return cropped/padded tensor with `sizes`"""
383        crop_to_sizes: Dict[AxisId, int] = {}
384        pad_to_sizes: Dict[AxisId, int] = {}
385        new_axes = dict(sizes)
386        for a, s_is in self.sizes.items():
387            a = AxisId(str(a))
388            _ = new_axes.pop(a, None)
389            if a not in sizes or sizes[a] == s_is:
390                pass
391            elif s_is > sizes[a]:
392                crop_to_sizes[a] = sizes[a]
393            else:
394                pad_to_sizes[a] = sizes[a]
395
396        tensor = self
397        if crop_to_sizes:
398            tensor = tensor.crop_to(crop_to_sizes, crop_where=crop_where)
399
400        if pad_to_sizes:
401            tensor = tensor.pad_to(pad_to_sizes, pad_where=pad_where, mode=pad_mode)
402
403        if new_axes:
404            tensor = tensor.expand_dims(new_axes)
405
406        return tensor
407
408    def transpose(
409        self,
410        axes: Sequence[AxisId],
411    ) -> Self:
412        """return a transposed tensor
413
414        Args:
415            axes: the desired tensor axes
416        """
417        # expand missing tensor axes
418        missing_axes = tuple(a for a in axes if a not in self.dims)
419        array = self._data
420        if missing_axes:
421            array = array.expand_dims(missing_axes)
422
423        # transpose to the correct axis order
424        return self.__class__.from_xarray(array.transpose(*axes))
425
426    @classmethod
427    def _interprete_array_wo_known_axes(cls, array: NDArray[Any]):
428        ndim = array.ndim
429        if ndim == 2:
430            current_axes = (
431                v0_5.SpaceInputAxis(id=AxisId("y"), size=array.shape[0]),
432                v0_5.SpaceInputAxis(id=AxisId("x"), size=array.shape[1]),
433            )
434        elif ndim == 3 and any(s <= 3 for s in array.shape):
435            current_axes = (
436                v0_5.ChannelAxis(
437                    channel_names=[
438                        v0_5.Identifier(f"channel{i}") for i in range(array.shape[0])
439                    ]
440                ),
441                v0_5.SpaceInputAxis(id=AxisId("y"), size=array.shape[1]),
442                v0_5.SpaceInputAxis(id=AxisId("x"), size=array.shape[2]),
443            )
444        elif ndim == 3:
445            current_axes = (
446                v0_5.SpaceInputAxis(id=AxisId("z"), size=array.shape[0]),
447                v0_5.SpaceInputAxis(id=AxisId("y"), size=array.shape[1]),
448                v0_5.SpaceInputAxis(id=AxisId("x"), size=array.shape[2]),
449            )
450        elif ndim == 4:
451            current_axes = (
452                v0_5.ChannelAxis(
453                    channel_names=[
454                        v0_5.Identifier(f"channel{i}") for i in range(array.shape[0])
455                    ]
456                ),
457                v0_5.SpaceInputAxis(id=AxisId("z"), size=array.shape[1]),
458                v0_5.SpaceInputAxis(id=AxisId("y"), size=array.shape[2]),
459                v0_5.SpaceInputAxis(id=AxisId("x"), size=array.shape[3]),
460            )
461        elif ndim == 5:
462            current_axes = (
463                v0_5.BatchAxis(),
464                v0_5.ChannelAxis(
465                    channel_names=[
466                        v0_5.Identifier(f"channel{i}") for i in range(array.shape[1])
467                    ]
468                ),
469                v0_5.SpaceInputAxis(id=AxisId("z"), size=array.shape[2]),
470                v0_5.SpaceInputAxis(id=AxisId("y"), size=array.shape[3]),
471                v0_5.SpaceInputAxis(id=AxisId("x"), size=array.shape[4]),
472            )
473        else:
474            raise ValueError(f"Could not guess an axis mapping for {array.shape}")
475
476        return cls(array, dims=tuple(a.id for a in current_axes))
477
478
479def _add_singletons(arr: NDArray[Any], axis_infos: Sequence[AxisInfo]):
480    if len(arr.shape) > len(axis_infos):
481        # remove singletons
482        for i, s in enumerate(arr.shape):
483            if s == 1:
484                arr = np.take(arr, 0, axis=i)
485                if len(arr.shape) == len(axis_infos):
486                    break
487
488    # add singletons if nececsary
489    for i, a in enumerate(axis_infos):
490        if len(arr.shape) >= len(axis_infos):
491            break
492
493        if a.maybe_singleton:
494            arr = np.expand_dims(arr, i)
495
496    return arr
497
498
499def _get_array_view(
500    original_array: NDArray[Any], axis_infos: Sequence[AxisInfo]
501) -> Optional[NDArray[Any]]:
502    perms = list(permutations(range(len(original_array.shape))))
503    perms.insert(1, perms.pop())  # try A and A.T first
504
505    for perm in perms:
506        view = original_array.transpose(perm)
507        view = _add_singletons(view, axis_infos)
508        if len(view.shape) != len(axis_infos):
509            return None
510
511        for s, a in zip(view.shape, axis_infos):
512            if s == 1 and not a.maybe_singleton:
513                break
514        else:
515            return view
516
517    return None