Source code for httomo.base_block

from typing import Tuple

import numpy as np

from httomo.block_interfaces import BlockData, BlockTransfer, generic_array
from httomo.runner.auxiliary_data import AuxiliaryData
from httomo.utils import gpu_enabled, make_3d_shape_from_array, xp


[docs] class BaseBlock(BlockData, BlockTransfer): """ Base block class providing default implementations for the data transferring/getting/setting behaviour needed for a block type to be processed by implementors of `MethodWrapper`. Ie, this class provides default implementations for the `BlockTransfer` and `BlockData` protocols. Note that the data indexing behaviour described in `DataIndexing` is not implemented in this class. If the default implementations for data transferring/getting/setting in this class are acceptable: inherit from `BaseBlock`, override where necessary, and implement `DataIndexing` in order to implement the `Block` protocol. """ def __init__(self, data: np.ndarray, aux_data: AuxiliaryData) -> None: self._data = data self._aux_data = aux_data def __dir__(self) -> list[str]: """Return only those properties that are relevant for the data""" return ["data", "angles", "angles_radians", "darks", "flats", "dark", "flat"] @property def data(self) -> generic_array: return self._data @data.setter def data(self, new_data: generic_array): self._data = new_data @property def aux_data(self) -> AuxiliaryData: return self._aux_data def _empty_aux_array(self): empty_shape = list(self._data.shape) return np.empty_like(self._data, shape=empty_shape) @property def angles(self) -> np.ndarray: return self._aux_data.get_angles() @angles.setter def angles(self, new_angles: np.ndarray): self._aux_data.set_angles(new_angles) @property def angles_radians(self) -> np.ndarray: return self.angles @angles_radians.setter def angles_radians(self, new_angles: np.ndarray): self.angles = new_angles @property def darks(self) -> generic_array: darks = self._aux_data.get_darks(self.is_gpu) if darks is None: darks = self._empty_aux_array() return darks @darks.setter def darks(self, darks: generic_array): self._aux_data.set_darks(darks) # alias @property def dark(self) -> generic_array: return self.darks @dark.setter def dark(self, darks: generic_array): self.darks = darks @property def flats(self) -> generic_array: flats = self._aux_data.get_flats(self.is_gpu) if flats is None: flats = self._empty_aux_array() return flats @flats.setter def flats(self, flats: generic_array): self._aux_data.set_flats(flats) # alias @property def flat(self) -> generic_array: return self.flats @flat.setter def flat(self, flats: generic_array): self.flats = flats @property def shape(self) -> Tuple[int, int, int]: """Shape of the data in this block""" return make_3d_shape_from_array(self._data)
[docs] def to_gpu(self): if not gpu_enabled: raise ValueError("no GPU available") self._data = xp.asarray(self.data, order="C")
[docs] def to_cpu(self): if not gpu_enabled: return self._data = xp.asnumpy(self.data, order="C")
@property def is_gpu(self) -> bool: return not self.is_cpu @property def is_cpu(self) -> bool: return getattr(self._data, "device", None) is None