Coverage for src/bioimageio/core/tensor.py: 82%
336 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-08 15:59 +0000
1from __future__ import annotations
3import collections.abc
4from itertools import permutations
5from typing import (
6 TYPE_CHECKING,
7 Any,
8 Callable,
9 Dict,
10 Iterator,
11 List,
12 Literal,
13 Mapping,
14 Optional,
15 Sequence,
16 Tuple,
17 Union,
18 cast,
19 get_args,
20)
22import numpy as np
23import pandas as pd
24import xarray as xr
25from loguru import logger
26from numpy.typing import DTypeLike, NDArray
27from typing_extensions import Self, assert_never
29from bioimageio.spec.model import v0_5
31from ._magic_tensor_ops import MagicTensorOpsMixin
32from .axis import AxisId, AxisInfo, AxisLike, PerAxis
33from .common import (
34 CropWhere,
35 DTypeStr,
36 PadMode,
37 PadWhere,
38 PadWidth,
39 PadWidthLike,
40 QuantileMethod,
41 SliceInfo,
42)
44if TYPE_CHECKING:
45 from numpy.typing import ArrayLike, NDArray
48_ScalarOrArray = Union["ArrayLike", np.generic, "NDArray[Any]"] # TODO: add "DaskArray"
51def _resolve_pad_mode(mode: PadMode):
52 constant_value = None
53 if isinstance(mode, str):
54 mode_name = mode
55 elif isinstance(mode, v0_5.ConstantPadding):
56 mode_name = mode.mode
57 constant_value = mode.value
58 elif isinstance(
59 mode, (v0_5.EdgePadding, v0_5.ReflectPadding, v0_5.SymmetricPadding)
60 ):
61 mode_name = mode.mode
62 else:
63 assert_never(mode)
65 return mode_name, constant_value
68# TODO: complete docstrings
69# TODO: in the long run---with improved typing in xarray---we should probably replace `Tensor` with xr.DataArray
70class Tensor(MagicTensorOpsMixin):
71 """A wrapper around an xr.DataArray for better integration with bioimageio.spec
72 and improved type annotations."""
74 _Compatible = Union["Tensor", xr.DataArray, _ScalarOrArray]
76 def __init__(
77 self,
78 array: Union[NDArray[Any], xr.DataArray],
79 dims: Sequence[Union[AxisId, AxisLike]],
80 ) -> None:
81 super().__init__()
82 axes = tuple(
83 a if isinstance(a, AxisId) else AxisInfo.create(a).id for a in dims
84 )
85 if isinstance(array, xr.DataArray):
86 self._data = array.transpose(*axes)
87 else:
88 self._data = xr.DataArray(array, dims=axes)
90 def __repr__(self) -> str:
91 return f"<Tensor {repr(self._data)}>"
93 def __array__(self, dtype: DTypeLike = None):
94 return np.asarray(self._data, dtype=dtype)
96 def __getitem__(
97 self,
98 key: Union[
99 SliceInfo,
100 slice,
101 int,
102 PerAxis[Union[SliceInfo, slice, int]],
103 Tensor,
104 xr.DataArray,
105 ],
106 ) -> Self:
107 if isinstance(key, SliceInfo):
108 key = slice(*key)
109 elif isinstance(key, collections.abc.Mapping):
110 key = {
111 a: s if isinstance(s, int) else s if isinstance(s, slice) else slice(*s)
112 for a, s in key.items()
113 }
114 elif isinstance(key, Tensor):
115 key = key._data
117 return self.__class__.from_xarray(self._data[key])
119 def __setitem__(
120 self,
121 key: Union[PerAxis[Union[SliceInfo, slice]], Tensor, xr.DataArray],
122 value: Union[Tensor, xr.DataArray, float, int],
123 ) -> None:
124 if isinstance(key, Tensor):
125 key = key._data
126 elif isinstance(key, xr.DataArray):
127 pass
128 else:
129 key = {a: s if isinstance(s, slice) else slice(*s) for a, s in key.items()}
131 if isinstance(value, Tensor):
132 value = value._data
134 self._data[key] = value
136 def __len__(self) -> int:
137 return len(self.data)
139 def _iter(self: Any) -> Iterator[Any]:
140 for n in range(len(self)):
141 yield self[n]
143 def __iter__(self: Any) -> Iterator[Any]:
144 if self.ndim == 0:
145 raise TypeError("iteration over a 0-d array")
146 return self._iter()
148 def _binary_op(
149 self,
150 other: _Compatible,
151 f: Callable[[Any, Any], Any],
152 reflexive: bool = False,
153 ) -> Self:
154 data = self._data._binary_op( # pyright: ignore[reportPrivateUsage]
155 (other._data if isinstance(other, Tensor) else other),
156 f,
157 reflexive,
158 )
159 return self.__class__.from_xarray(data)
161 def _inplace_binary_op(
162 self,
163 other: _Compatible,
164 f: Callable[[Any, Any], Any],
165 ) -> Self:
166 _ = self._data._inplace_binary_op( # pyright: ignore[reportPrivateUsage]
167 (
168 other_d
169 if (other_d := getattr(other, "data")) is not None
170 and isinstance(
171 other_d,
172 xr.DataArray,
173 )
174 else other
175 ),
176 f,
177 )
178 return self
180 def _unary_op(self, f: Callable[[Any], Any], *args: Any, **kwargs: Any) -> Self:
181 data = self._data._unary_op( # pyright: ignore[reportPrivateUsage]
182 f, *args, **kwargs
183 )
184 return self.__class__.from_xarray(data)
186 @classmethod
187 def from_xarray(cls, data_array: xr.DataArray) -> Self:
188 """create a `Tensor` from an xarray data array
190 note for internal use: this factory method is round-trip save
191 for any `Tensor`'s `data` property (an xarray.DataArray).
192 """
193 return cls(array=data_array, dims=tuple(AxisId(d) for d in data_array.dims))
195 @classmethod
196 def from_numpy(
197 cls,
198 array: NDArray[Any],
199 *,
200 dims: Optional[Union[AxisLike, Sequence[AxisLike]]],
201 ) -> Tensor:
202 """create a `Tensor` from a numpy array
204 Args:
205 array: the nd numpy array
206 dims: A description of the array's axes.
207 If None axes are guessed (which might fail and raise a ValueError.)
208 If dims do not match array shape, permutations and singleton dimensions are tried to find a match.
209 Raises:
210 ValueError: if `dims` is None and dims guessing fails.
211 """
213 if dims is None:
214 return cls._interprete_array_wo_known_axes(array)
215 elif isinstance(dims, collections.abc.Sequence):
216 dim_seq = list(dims)
217 else:
218 dim_seq = [dims]
220 axis_infos = [AxisInfo.create(a) for a in dim_seq]
221 original_shape = tuple(array.shape)
223 successful_view = _get_array_view(array, axis_infos)
224 if successful_view is None:
225 raise ValueError(
226 f"Array shape {original_shape} does not map to axes {dims}"
227 )
229 return Tensor(successful_view, dims=tuple(a.id for a in axis_infos))
231 @property
232 def data(self):
233 return self._data
235 @property
236 def dims(self): # TODO: rename to `axes`?
237 """Tuple of dimension names associated with this tensor."""
238 return cast(Tuple[AxisId, ...], self._data.dims)
240 @property
241 def dtype(self) -> DTypeStr:
242 dt = str(self.data.dtype) # pyright: ignore[reportUnknownArgumentType]
243 assert dt in get_args(DTypeStr)
244 return dt # pyright: ignore[reportReturnType]
246 @property
247 def ndim(self):
248 """Number of tensor dimensions."""
249 return self._data.ndim
251 @property
252 def shape(self):
253 """Tuple of tensor axes lengths"""
254 return self._data.shape
256 @property
257 def shape_tuple(self):
258 """Tuple of tensor axes lengths"""
259 return self._data.shape
261 @property
262 def size(self):
263 """Number of elements in the tensor.
265 Equal to math.prod(tensor.shape), i.e., the product of the tensors’ dimensions.
266 """
267 return self._data.size
269 @property
270 def sizes(self):
271 """Ordered, immutable mapping from axis ids to axis lengths."""
272 return cast(Mapping[AxisId, int], self.data.sizes)
274 @property
275 def tagged_shape(self):
276 """(alias for `sizes`) Ordered, immutable mapping from axis ids to lengths."""
277 return self.sizes
279 def to_numpy(self) -> NDArray[Any]:
280 """Return the data of this tensor as a numpy array."""
281 return self.data.to_numpy() # pyright: ignore[reportUnknownVariableType]
283 def argmax(self) -> Mapping[AxisId, int]:
284 ret = self._data.argmax(...)
285 assert isinstance(ret, dict)
286 return {cast(AxisId, k): cast(int, v.item()) for k, v in ret.items()}
288 def astype(self, dtype: DTypeStr, *, copy: bool = False):
289 """Return tensor cast to `dtype`
291 note: if dtype is already satisfied copy if `copy`"""
292 return self.__class__.from_xarray(self._data.astype(dtype, copy=copy))
294 def clip(self, min: Optional[float] = None, max: Optional[float] = None):
295 """Return a tensor whose values are limited to [min, max].
296 At least one of max or min must be given."""
297 return self.__class__.from_xarray(self._data.clip(min, max))
299 def crop_to(
300 self,
301 sizes: PerAxis[int],
302 crop_where: Union[
303 CropWhere,
304 PerAxis[CropWhere],
305 ] = "left_and_right",
306 ) -> Self:
307 """crop to match `sizes`"""
308 if isinstance(crop_where, str):
309 crop_axis_where: PerAxis[CropWhere] = {a: crop_where for a in self.dims}
310 else:
311 crop_axis_where = crop_where
313 slices: Dict[AxisId, SliceInfo] = {}
315 for a, s_is in self.sizes.items():
316 if a not in sizes or sizes[a] == s_is:
317 pass
318 elif sizes[a] > s_is:
319 logger.warning(
320 "Cannot crop axis {} of size {} to larger size {}",
321 a,
322 s_is,
323 sizes[a],
324 )
325 elif a not in crop_axis_where:
326 raise ValueError(
327 f"Don't know where to crop axis {a}, `crop_where`={crop_where}"
328 )
329 else:
330 crop_this_axis_where = crop_axis_where[a]
331 if crop_this_axis_where == "left":
332 slices[a] = SliceInfo(s_is - sizes[a], s_is)
333 elif crop_this_axis_where == "right":
334 slices[a] = SliceInfo(0, sizes[a])
335 elif crop_this_axis_where == "left_and_right":
336 slices[a] = SliceInfo(
337 start := (s_is - sizes[a]) // 2, sizes[a] + start
338 )
339 else:
340 assert_never(crop_this_axis_where)
342 return self[slices]
344 def expand_dims(self, dims: Union[Sequence[AxisId], PerAxis[int]]) -> Self:
345 return self.__class__.from_xarray(self._data.expand_dims(dims=dims))
347 def item(
348 self,
349 key: Union[
350 None, SliceInfo, slice, int, PerAxis[Union[SliceInfo, slice, int]]
351 ] = None,
352 ):
353 """Copy a tensor element to a standard Python scalar and return it."""
354 if key is None:
355 ret = self._data.item()
356 else:
357 ret = self[key]._data.item()
359 assert isinstance(ret, (bool, float, int))
360 return ret
362 def mean(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
363 return self.__class__.from_xarray(self._data.mean(dim=dim))
365 def pad(
366 self,
367 pad_width: PerAxis[PadWidthLike],
368 mode: PadMode = "symmetric",
369 ) -> Self:
370 pad_width = {a: PadWidth.create(p) for a, p in pad_width.items()}
371 mode_name, constant_value = _resolve_pad_mode(mode)
372 return self.__class__.from_xarray(
373 self._data.pad(
374 pad_width=pad_width, mode=mode_name, constant_values=constant_value
375 )
376 )
378 def pad_to(
379 self,
380 sizes: PerAxis[int],
381 pad_where: Union[PadWhere, PerAxis[PadWhere]] = "left_and_right",
382 mode: PadMode = "symmetric",
383 ) -> Self:
384 """pad `tensor` to match `sizes`"""
385 if isinstance(pad_where, str):
386 pad_axis_where: PerAxis[PadWhere] = {a: pad_where for a in self.dims}
387 else:
388 pad_axis_where = pad_where
390 pad_width: Dict[AxisId, PadWidth] = {}
391 for a, s_is in self.sizes.items():
392 if a not in sizes or sizes[a] == s_is:
393 pad_width[a] = PadWidth(0, 0)
394 elif s_is > sizes[a]:
395 pad_width[a] = PadWidth(0, 0)
396 logger.warning(
397 "Cannot pad axis {} of size {} to smaller size {}",
398 a,
399 s_is,
400 sizes[a],
401 )
402 elif a not in pad_axis_where:
403 raise ValueError(
404 f"Don't know where to pad axis {a}, `pad_where`={pad_where}"
405 )
406 else:
407 pad_this_axis_where = pad_axis_where[a]
408 d = sizes[a] - s_is
409 if pad_this_axis_where == "left":
410 pad_width[a] = PadWidth(d, 0)
411 elif pad_this_axis_where == "right":
412 pad_width[a] = PadWidth(0, d)
413 elif pad_this_axis_where == "left_and_right":
414 pad_width[a] = PadWidth(left := d // 2, d - left)
415 else:
416 assert_never(pad_this_axis_where)
418 return self.pad(pad_width, mode)
420 def quantile(
421 self,
422 q: Union[float, Sequence[float]],
423 dim: Optional[Union[AxisId, Sequence[AxisId]]] = None,
424 method: QuantileMethod = "linear",
425 ) -> Self:
426 assert (
427 isinstance(q, (float, int))
428 and q >= 0.0
429 or not isinstance(q, (float, int))
430 and all(qq >= 0.0 for qq in q)
431 )
432 assert (
433 isinstance(q, (float, int))
434 and q <= 1.0
435 or not isinstance(q, (float, int))
436 and all(qq <= 1.0 for qq in q)
437 )
438 assert dim is None or (
439 (quantile_dim := AxisId("quantile")) != dim and quantile_dim not in set(dim)
440 )
441 return self.__class__.from_xarray(
442 self._data.quantile(q, dim=dim, method=method)
443 )
445 def resize_to(
446 self,
447 sizes: PerAxis[int],
448 *,
449 pad_where: Union[
450 PadWhere,
451 PerAxis[PadWhere],
452 ] = "left_and_right",
453 crop_where: Union[
454 CropWhere,
455 PerAxis[CropWhere],
456 ] = "left_and_right",
457 pad_mode: PadMode = "symmetric",
458 ):
459 """return cropped/padded tensor with `sizes`"""
460 crop_to_sizes: Dict[AxisId, int] = {}
461 pad_to_sizes: Dict[AxisId, int] = {}
462 new_axes = dict(sizes)
463 for a, s_is in self.sizes.items():
464 a = AxisId(str(a))
465 _ = new_axes.pop(a, None)
466 if a not in sizes or sizes[a] == s_is:
467 pass
468 elif s_is > sizes[a]:
469 crop_to_sizes[a] = sizes[a]
470 else:
471 pad_to_sizes[a] = sizes[a]
473 tensor = self
474 if crop_to_sizes:
475 tensor = tensor.crop_to(crop_to_sizes, crop_where=crop_where)
477 if pad_to_sizes:
478 tensor = tensor.pad_to(pad_to_sizes, pad_where=pad_where, mode=pad_mode)
480 if new_axes:
481 tensor = tensor.expand_dims(new_axes)
483 return tensor
485 def std(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
486 return self.__class__.from_xarray(self._data.std(dim=dim))
488 def sum(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
489 """Reduce this Tensor's data by applying sum along some dimension(s)."""
490 return self.__class__.from_xarray(self._data.sum(dim=dim))
492 def assign_batch_multi_index(self, multi_index: "pd.MultiIndex") -> Self:
493 """Set the batch multi-index for this tensor.
495 Args:
496 multi_index: The multi-index to set.
497 """
498 if AxisId("batch") not in self.dims:
499 raise ValueError(
500 "Cannot set batch multi-index on a tensor without a 'batch' axis."
501 )
503 return self.__class__.from_xarray(
504 self._data.assign_coords({AxisId("batch"): multi_index})
505 )
507 def unstack_batch_multi_index(
508 self, *, errors: Literal["raise", "ignore"] = "raise"
509 ) -> Self:
510 """Unstack the batch multi-index of this tensor.
512 Returns:
513 A new tensor with the batch multi-index unstacked into separate axes.
514 """
515 if AxisId("batch") not in self.dims:
516 if errors == "raise":
517 raise ValueError(
518 "Cannot unstack batch multi-index on a tensor without a 'batch' axis."
519 )
520 elif errors == "ignore":
521 return self
522 else:
523 assert_never(errors)
525 if not isinstance(self._data.indexes.get(AxisId("batch")), pd.MultiIndex):
526 if errors == "raise":
527 raise ValueError(
528 "Cannot unstack batch multi-index on a tensor whose 'batch' axis does not have a MultiIndex."
529 )
530 elif errors == "ignore":
531 return self
532 else:
533 assert_never(errors)
535 old_dims = self.dims
536 array = self._data.unstack(AxisId("batch"))
537 added_dims = [AxisId(d) for d in array.dims if d not in self._data.dims]
539 # restore expected axis order, replace batch dim with added dims
540 new_dims: List[AxisId] = []
541 for d in old_dims:
542 if d in array.dims:
543 new_dims.append(d)
544 elif d == AxisId("batch"):
545 new_dims.extend(added_dims)
546 else:
547 raise ValueError(f"Expected axis {d} not found in unstacked array.")
549 array = array.transpose(*new_dims)
550 if AxisId("original_batch") in array.dims:
551 array = array.rename({AxisId("original_batch"): AxisId("batch")})
553 return self.__class__.from_xarray(array)
555 def transpose(
556 self,
557 axes: Sequence[AxisId],
558 *,
559 extra_dims: Literal[
560 "raise", "squeeze", "stack", "squeeze_or_stack"
561 ] = "squeeze",
562 missing_dims: Literal[
563 "raise", "expand", "unstack", "unstack_or_expand"
564 ] = "unstack_or_expand",
565 ) -> Self:
566 """Return a transposed tensor, missing axes are expanded (if `unstack_missing_dims_from_batch` is False) or unstacked from batch (if `unstack_missing_dims_from_batch` is True), extra axes are stacked to batch (if `stack_extra_dims_to_batch` is True). Additional axes raise (if `stack_extra_dims_to_batch` is True).
568 Args:
569 axes: The desired tensor axes
570 extra_dims:
571 Extra dimensions are any dimensions in the tensor that are not specified in `axes`.
572 If "raise", any extra dimensions will raise an error.
573 If "squeeze", any extra singleton dimensions will be squeezed, non-singleton dimensions will raise an error.
574 If "stack", any extra dimensions will be stacked to the batch dimension. Such a stacked batch dimension then has a multi-index that can be unstacked using `Tensor.unstack_batch_multi_index()`.
575 If "squeeze_or_stack", any extra singleton dimensions will be squeezed, non-singleton dimensions will be stacked to the batch dimension.
576 missing_dims:
577 Missing dimensions are any dimensions specified in `axes` that are not present in the tensor.
578 If "raise", any missing dimensions will raise an error.
579 If "expand", any missing dimensions will be added as singleton dimensions.
580 If "unstack", any missing dimensions will be unstacked from the batch dimension. For this option a batch dimension with a multi-index must be present from previous stacking operations or assigned by `Tensor.assign_batch_multi_index()`.
581 If "unstack_or_expand", any missing dimensions will be unstacked from the batch dimension if it has a multi-index, otherwise they will be added as singleton dimensions.
582 """
583 array = self._data
585 unhandled_missing_dims = [a for a in axes if a not in array.dims]
586 if unhandled_missing_dims and missing_dims == "raise":
587 raise ValueError(f"Found missing dimensions {unhandled_missing_dims}.")
589 unstack_error = None
590 if unhandled_missing_dims and missing_dims in ("unstack", "unstack_or_expand"):
591 lets_unstack = AxisId("batch") in array.dims
592 if not lets_unstack:
593 unstack_error = f"Missing dimensions {unhandled_missing_dims} found, but 'batch' axis is not in the tensor. Cannot unstack missing dimensions from batch."
594 if missing_dims == "unstack":
595 raise ValueError(unstack_error)
597 if lets_unstack and not isinstance(
598 array.indexes.get(AxisId("batch")), pd.MultiIndex
599 ):
600 lets_unstack = False
601 unstack_error = f"Missing dimensions {unhandled_missing_dims} found, but 'batch' axis does not have a MultiIndex. Cannot unstack missing dimensions from non-multi-index batch."
602 if missing_dims == "unstack":
603 raise ValueError(unstack_error)
604 else:
605 lets_unstack = False
607 if lets_unstack:
608 array = array.unstack(AxisId("batch"))
610 if AxisId("original_batch") in array.dims:
611 if AxisId("batch") in axes:
612 array = array.rename({AxisId("original_batch"): AxisId("batch")})
613 else:
614 array = array.squeeze(AxisId("original_batch"))
616 unhandled_missing_dims = [a for a in axes if a not in array.dims]
618 if unhandled_missing_dims and missing_dims in ("expand", "unstack_or_expand"):
619 array = array.expand_dims(unhandled_missing_dims)
620 unhandled_missing_dims = []
622 if unhandled_missing_dims:
623 if unstack_error is not None:
624 raise ValueError(unstack_error)
626 raise ValueError(f"Missing dimensions {unhandled_missing_dims}.")
628 unhandled_extra_dims = [a for a in array.dims if a not in axes]
630 if unhandled_extra_dims and extra_dims == "raise":
631 raise ValueError(f"Found extra dimensions {unhandled_extra_dims}.")
633 if unhandled_extra_dims and extra_dims in ("squeeze", "squeeze_or_stack"):
634 for d in list(unhandled_extra_dims):
635 if array.sizes[d] == 1:
636 array = array.squeeze(d)
637 unhandled_extra_dims.remove(d)
638 elif extra_dims == "squeeze":
639 raise ValueError(
640 f"Extra dimension {d} found but stack_extra_dims_to_batch is False and the dimension is not a singleton."
641 )
643 if unhandled_extra_dims and extra_dims in ("stack", "squeeze_or_stack"):
644 if AxisId("batch") not in axes:
645 raise ValueError(
646 f"Extra dimensions {unhandled_extra_dims} found but 'batch' axis is not in the desired axes {axes}."
647 + " Cannot stack extra dimensions to batch."
648 )
650 if AxisId("batch") in array.dims:
651 array = array.rename({AxisId("batch"): AxisId("original_batch")})
652 unhandled_extra_dims.insert(0, AxisId("original_batch"))
654 array = array.stack({AxisId("batch"): unhandled_extra_dims})
655 unhandled_extra_dims = []
657 if unhandled_extra_dims:
658 raise ValueError(
659 f"Non-singleton extra dimensions {unhandled_extra_dims} found, but `extra_dims` not in ('stack', 'squeeze_or_stack')."
660 )
662 # transpose to the correct axis order
663 return self.__class__.from_xarray(array.transpose(*axes))
665 def var(self, dim: Optional[Union[AxisId, Sequence[AxisId]]] = None) -> Self:
666 return self.__class__.from_xarray(self._data.var(dim=dim))
668 @classmethod
669 def _interprete_array_wo_known_axes(cls, array: NDArray[Any]):
670 ndim = array.ndim
671 if ndim == 2:
672 current_axes = (
673 v0_5.SpaceInputAxis(id=v0_5.AxisId("y"), size=array.shape[0]),
674 v0_5.SpaceInputAxis(id=v0_5.AxisId("x"), size=array.shape[1]),
675 )
676 elif ndim == 3 and any(s <= 3 for s in array.shape):
677 current_axes = (
678 v0_5.ChannelAxis(
679 channel_names=[
680 v0_5.Identifier(f"channel{i}") for i in range(array.shape[0])
681 ]
682 ),
683 v0_5.SpaceInputAxis(id=v0_5.AxisId("y"), size=array.shape[1]),
684 v0_5.SpaceInputAxis(id=v0_5.AxisId("x"), size=array.shape[2]),
685 )
686 elif ndim == 3:
687 current_axes = (
688 v0_5.SpaceInputAxis(id=v0_5.AxisId("z"), size=array.shape[0]),
689 v0_5.SpaceInputAxis(id=v0_5.AxisId("y"), size=array.shape[1]),
690 v0_5.SpaceInputAxis(id=v0_5.AxisId("x"), size=array.shape[2]),
691 )
692 elif ndim == 4:
693 current_axes = (
694 v0_5.ChannelAxis(
695 channel_names=[
696 v0_5.Identifier(f"channel{i}") for i in range(array.shape[0])
697 ]
698 ),
699 v0_5.SpaceInputAxis(id=v0_5.AxisId("z"), size=array.shape[1]),
700 v0_5.SpaceInputAxis(id=v0_5.AxisId("y"), size=array.shape[2]),
701 v0_5.SpaceInputAxis(id=v0_5.AxisId("x"), size=array.shape[3]),
702 )
703 elif ndim == 5:
704 current_axes = (
705 v0_5.BatchAxis(),
706 v0_5.ChannelAxis(
707 channel_names=[
708 v0_5.Identifier(f"channel{i}") for i in range(array.shape[1])
709 ]
710 ),
711 v0_5.SpaceInputAxis(id=v0_5.AxisId("z"), size=array.shape[2]),
712 v0_5.SpaceInputAxis(id=v0_5.AxisId("y"), size=array.shape[3]),
713 v0_5.SpaceInputAxis(id=v0_5.AxisId("x"), size=array.shape[4]),
714 )
715 else:
716 raise ValueError(f"Could not guess an axis mapping for {array.shape}")
718 return cls(array, dims=tuple(a.id for a in current_axes))
721def _add_singletons(arr: NDArray[Any], axis_infos: Sequence[AxisInfo]):
722 if len(arr.shape) > len(axis_infos):
723 # remove singletons
724 for i, s in enumerate(arr.shape):
725 if s == 1:
726 arr = np.take(arr, 0, axis=i)
727 if len(arr.shape) == len(axis_infos):
728 break
730 # add singletons if nececsary
731 for i, a in enumerate(axis_infos):
732 if len(arr.shape) >= len(axis_infos):
733 break
735 if a.size.min == 1:
736 arr = np.expand_dims(arr, i)
738 return arr
741def _get_array_view(
742 original_array: NDArray[Any], axis_infos: Sequence[AxisInfo]
743) -> Optional[NDArray[Any]]:
744 perms = list(permutations(range(len(original_array.shape))))
746 for perm in perms:
747 view = original_array.transpose(perm)
748 view = _add_singletons(view, axis_infos)
749 if len(view.shape) != len(axis_infos):
750 return None
752 for s, a in zip(view.shape, axis_infos):
753 if (
754 s < a.size.min
755 or (a.size.max is not None and s > a.size.max)
756 or (a.size.step is not None and (s - a.size.min) % a.size.step != 0)
757 ):
758 break
759 else:
760 return view
762 return None