from typing import List, Literal, Optional, Tuple
import numpy as np
from httomo.base_block import BaseBlock, generic_array
from httomo.block_interfaces import BlockIndexing
from httomo.runner.auxiliary_data import AuxiliaryData
from httomo.utils import make_3d_shape_from_array, make_3d_shape_from_shape
[docs]
class DataSetBlock(BaseBlock, BlockIndexing):
"""
Data storage type for block processing in high throughput runs
"""
def __init__(
self,
data: np.ndarray,
aux_data: AuxiliaryData,
slicing_dim: Literal[0, 1, 2] = 0,
block_start: int = 0,
chunk_start: int = 0,
global_shape: Optional[Tuple[int, int, int]] = None,
chunk_shape: Optional[Tuple[int, int, int]] = None,
padding: Tuple[int, int] = (0, 0),
):
"""Constructs a data block for processing in the pipeline in high throughput runs.
Parameters
----------
data: ndarray
A numpy or cupy array, 3D, holding the data represented by this block.
aux_data: AuxiliaryData
Object handling the flats, darks, and angles arrays
slicing_dim: Literal[0, 1, 2]
The slicing dimension in the global data that this block represents as slice of.
This is to facilitate parallel processing - data is sliced in one of the 3 dimensions.
block_start: int
The index in slicing dimensions within the chunk that this block starts at. It is relative
to the start of the chunk.
chunk_start: int
The index in slicing dimension within the global data that the underlying chunk starts.
A chunk is a unit of the global data that is handled by a single MPI process, while a block
might be a smaller part than the chunk.
global_shape: Optional[Tuple[int, int, int]]
The shape of the global data across all processes. If not given, it assumes this block
represents the full global data (no slicing done).
chunk_shape: Optional[Tuple[int, int, int]]
The shape of the chunk that this block belongs to. If not given, it assumes this block
spans the full chunk.
padding: Tuple[int, int]
Padding information - holds the number of padded slices before and after the core area of the block,
in slicing dimension. If not given, no padding is assumed.
Note that the padding information should be added to the data's shape, i.e. block_start, chunk_start,
chunk_shape, and the data's shape includes the padded slices. And therefore block_start or chunk_start
may have negative values of up to -padding[0]. The global_shape is not adapted for padding.
"""
super().__init__(data, aux_data)
self._slicing_dim = slicing_dim
self._block_start = block_start
self._chunk_start = chunk_start
self._padding = padding
if global_shape is None:
global_shape_t = list(data.shape)
global_shape_t[slicing_dim] -= padding[0] + padding[1]
self._global_shape = make_3d_shape_from_shape(global_shape_t)
else:
self._global_shape = global_shape
if chunk_shape is None:
self._chunk_shape = make_3d_shape_from_array(data)
else:
self._chunk_shape = chunk_shape
chunk_index = [0, 0, 0]
chunk_index[slicing_dim] += block_start
self._chunk_index = make_3d_shape_from_shape(chunk_index)
global_index = [0, 0, 0]
global_index[slicing_dim] += chunk_start + block_start + padding[0]
self._global_index = make_3d_shape_from_shape(global_index)
self._check_inconsistencies()
def _check_inconsistencies(self):
if self.padding[0] < 0 or self.padding[1] < 0:
raise ValueError("padding values cannot be negative")
if self.chunk_index[self.slicing_dim] + self._padding[0] < 0:
raise ValueError("block start index must be >= 0")
if (
self.chunk_index[self.slicing_dim]
+ self.shape[self.slicing_dim]
- self._padding[1]
> self.chunk_shape[self.slicing_dim]
):
raise ValueError("block spans beyond the chunk's boundaries")
if self.global_index[self.slicing_dim] + self._padding[0] < 0:
raise ValueError("chunk start index must be >= 0")
if (
self.global_index[self.slicing_dim]
+ self.shape[self.slicing_dim]
- self._padding[1]
> self.global_shape[self.slicing_dim]
):
raise ValueError("chunk spans beyond the global data boundaries")
if any(
self.chunk_shape[i] > self.global_shape[i]
for i in range(3)
if i != self.slicing_dim
):
raise ValueError(
"chunk shape is larger than the global shape in non-slicing dimensions"
)
if (
self.chunk_shape[self.slicing_dim] - self.padding[0] - self.padding[1]
> self.global_shape[self.slicing_dim]
):
raise ValueError(
"chunk shape is larger than the global shape in slicing dimension"
)
if any(self.shape[i] > self.chunk_shape[i] for i in range(3)):
raise ValueError("block shape is larger than the chunk shape")
if any(
self.shape[i] != self.global_shape[i]
for i in range(3)
if i != self.slicing_dim
):
raise ValueError(
"block shape inconsistent with non-slicing dims of global shape"
)
assert not any(
self.chunk_shape[i] != self.global_shape[i]
for i in range(3)
if i != self.slicing_dim
)
@property
def chunk_index(self) -> Tuple[int, int, int]:
"""The index of this block within the chunk handled by the current process"""
return self._chunk_index
@property
def chunk_shape(self) -> Tuple[int, int, int]:
"""Shape of the full chunk handled by the current process"""
return self._chunk_shape
@property
def global_index(self) -> Tuple[int, int, int]:
"""The index of this block within the global data across all processes"""
return self._global_index
@property
def global_shape(self) -> Tuple[int, int, int]:
"""Shape of the global data across all processes"""
return self._global_shape
@property
def is_last_in_chunk(self) -> bool:
"""Check if the current dataset is the final one for the chunk handled by the current process"""
return (
self.chunk_index[self._slicing_dim] + self.shape[self._slicing_dim]
== self.chunk_shape[self._slicing_dim] - self.padding[0]
)
@property
def slicing_dim(self) -> Literal[0, 1, 2]:
return self._slicing_dim
def _empty_aux_array(self):
empty_shape = list(self._data.shape)
empty_shape[self.slicing_dim] = 0
return np.empty_like(self._data, shape=empty_shape)
@property
def data(self) -> generic_array:
return super().data
@data.setter
def data(self, new_data: generic_array):
global_shape = list(self._global_shape)
chunk_shape = list(self._chunk_shape)
for i in range(3):
if i != self.slicing_dim:
global_shape[i] = new_data.shape[i]
chunk_shape[i] = new_data.shape[i]
elif self._data.shape[i] != new_data.shape[i]:
raise ValueError("shape mismatch in slicing dimension")
self._data = new_data
self._global_shape = make_3d_shape_from_shape(global_shape)
self._chunk_shape = make_3d_shape_from_shape(chunk_shape)
@data.deleter
def data(self):
del self._data
del self._global_shape
del self._chunk_shape
@property
def is_padded(self) -> bool:
return self._padding != (0, 0)
@property
def padding(self) -> Tuple[int, int]:
return self._padding
@property
def shape_unpadded(self) -> Tuple[int, int, int]:
return self._correct_shape_for_padding(self.shape)
@property
def chunk_index_unpadded(self) -> Tuple[int, int, int]:
return self._correct_index_for_padding(self.chunk_index)
@property
def chunk_shape_unpadded(self) -> Tuple[int, int, int]:
return self._correct_shape_for_padding(self.chunk_shape)
@property
def global_index_unpadded(self) -> Tuple[int, int, int]:
return self._correct_index_for_padding(self.global_index)
def _correct_shape_for_padding(
self, shape: Tuple[int, int, int]
) -> Tuple[int, int, int]:
if not self.padding:
return shape
shp = list(shape)
shp[self.slicing_dim] -= self.padding[0] + self.padding[1]
return make_3d_shape_from_shape(shp)
def _correct_index_for_padding(
self, index: Tuple[int, int, int]
) -> Tuple[int, int, int]:
if not self.padding:
return index
idx = list(index)
idx[self.slicing_dim] += self.padding[0]
return make_3d_shape_from_shape(idx)
@property
def data_unpadded(self) -> generic_array:
if not self.padding:
return self.data
d = self.data
slices = [slice(None), slice(None), slice(None)]
slices[self.slicing_dim] = slice(
self.padding[0], d.shape[self.slicing_dim] - self.padding[1]
)
return d[slices[0], slices[1], slices[2]]
