From b2857dbb66ca2f88506a3ecd8b530e1813501d74 Mon Sep 17 00:00:00 2001 From: ashmeigh Date: Tue, 16 Jul 2024 18:59:42 +0100 Subject: [PATCH] Added type annobnations to core Signed-off-by: ashmeigh --- mantidimaging/core/gpu/utility.py | 28 +++++++-------- .../core/io/instrument_log_implmentations.py | 4 +-- mantidimaging/core/io/saver.py | 21 ++++++----- mantidimaging/core/net/help_pages.py | 4 +-- mantidimaging/core/parallel/shared.py | 2 +- mantidimaging/core/parallel/utility.py | 6 ++-- mantidimaging/core/reconstruct/astra_recon.py | 4 +-- mantidimaging/core/reconstruct/base_recon.py | 2 +- .../core/reconstruct/tomopy_recon.py | 4 +-- mantidimaging/core/rotation/data_model.py | 36 +++++++++---------- .../core/rotation/polyfit_correlation.py | 6 ++-- 11 files changed, 60 insertions(+), 57 deletions(-) diff --git a/mantidimaging/core/gpu/utility.py b/mantidimaging/core/gpu/utility.py index bc0e74af5bb..12b34f20c7b 100644 --- a/mantidimaging/core/gpu/utility.py +++ b/mantidimaging/core/gpu/utility.py @@ -32,14 +32,14 @@ } -def _cupy_on_system(): +def _cupy_on_system() -> bool: """ :return: True if cupy is installed on the system, False otherwise. """ return not CUPY_NOT_IMPORTED -def _cupy_installed_correctly(): +def _cupy_installed_correctly() -> bool: """ :return: True if cupy is able to run on the system, False otherwise. """ @@ -64,14 +64,14 @@ def _cupy_installed_correctly(): return False -def gpu_available(): +def gpu_available() -> bool: """ :return: True if cupy is installed AND working, False otherwise. """ return _cupy_on_system() and _cupy_installed_correctly() -def _load_cuda_kernel(dtype): +def _load_cuda_kernel(dtype) -> str: """ Loads the CUDA kernel so that cupy can act as a mediator. Replaces instances of 'float' with 'double' if the dtype is float64. @@ -86,7 +86,7 @@ def _load_cuda_kernel(dtype): return cuda_kernel -def _free_memory_pool(arrays=None): +def _free_memory_pool(arrays=None) -> None: """ Delete any given GPU arrays and instruct the memory pool to free unused blocks. """ @@ -95,7 +95,7 @@ def _free_memory_pool(arrays=None): mempool.free_all_blocks() -def _create_pinned_memory(cpu_array): +def _create_pinned_memory(cpu_array: np.ndarray) -> np.ndarray: """ Use pinned memory in order to store a numpy array on the GPU. :param cpu_array: The numpy array to be transferred to the GPU. @@ -107,7 +107,7 @@ def _create_pinned_memory(cpu_array): return src -def _send_single_array_to_gpu(cpu_array, stream): +def _send_single_array_to_gpu(cpu_array: np.ndarray, stream: cp.cuda.Stream) -> cp.ndarray: """ Sends a single array to the GPU using pinned memory and a stream. :param cpu_array: The numpy array to be transferred to the GPU. @@ -120,7 +120,7 @@ def _send_single_array_to_gpu(cpu_array, stream): return gpu_array -def _send_arrays_to_gpu_with_pinned_memory(cpu_arrays, streams): +def _send_arrays_to_gpu_with_pinned_memory(cpu_arrays, streams) -> list[cp.ndarray]: """ Transfer the arrays to the GPU using pinned memory. Raises an error if the GPU runs out of memory. :param cpu_arrays: A list of numpy arrays to be transferred to the GPU. @@ -145,7 +145,7 @@ def _send_arrays_to_gpu_with_pinned_memory(cpu_arrays, streams): return [] -def _create_block_and_grid_args(data): +def _create_block_and_grid_args(data: cp.ndarray): """ Create the block and grid arguments that are passed to the cupy. These determine how the array is broken up. @@ -158,7 +158,7 @@ def _create_block_and_grid_args(data): return block_size, grid_size -def _create_padded_array(data, filter_size, scipy_mode): +def _create_padded_array(data: np.ndarray, filter_size: int, scipy_mode: str) -> np.ndarray: """ Creates the padded array on the CPU for the median filter. :param data: The data array to be padded. @@ -171,7 +171,7 @@ def _create_padded_array(data, filter_size, scipy_mode): return np.pad(data, pad_width=((pad_size, pad_size), (pad_size, pad_size)), mode=EQUIVALENT_PAD_MODE[scipy_mode]) -def _replace_gpu_array_contents(gpu_array, cpu_array, stream): +def _replace_gpu_array_contents(gpu_array: cp.ndarray, cpu_array: np.ndarray, stream: cp.cuda.Stream) -> None: """ Overwrites the contents of an existing GPU array with a given CPU array. :param gpu_array: The GPU array to be overwritten. @@ -181,7 +181,7 @@ def _replace_gpu_array_contents(gpu_array, cpu_array, stream): gpu_array.set(cpu_array, stream) -def _get_padding_value(filter_size): +def _get_padding_value(filter_size: int) -> int: """ Determine the padding value by using the filter size. :param filter_size: The filter size. @@ -202,7 +202,7 @@ def __init__(self, dtype): # Warm up the CUDA functions self._warm_up(dtype) - def _warm_up(self, dtype): + def _warm_up(self, dtype) -> None: """ Runs the median filter on a small test array in order to allow it to compile then deleted the GPU arrays. :param dtype: The data type of the input array. @@ -219,7 +219,7 @@ def _warm_up(self, dtype): # Clear the test arrays _free_memory_pool([test_data, test_padding]) - def _cuda_single_image_median_filter(self, input_data, padded_data, filter_size, grid_size, block_size): + def _cuda_single_image_median_filter(self, input_data, padded_data, filter_size, grid_size, block_size) -> None: """ Run the median filter on a single 2D image using CUDA. :param input_data: A 2D GPU data array. diff --git a/mantidimaging/core/io/instrument_log_implmentations.py b/mantidimaging/core/io/instrument_log_implmentations.py index 248f199a8dd..1400d6460d1 100644 --- a/mantidimaging/core/io/instrument_log_implmentations.py +++ b/mantidimaging/core/io/instrument_log_implmentations.py @@ -100,7 +100,7 @@ def read_imat_date(time_stamp: str) -> datetime: locale.setlocale(locale.LC_TIME, lc) @staticmethod - def _has_imat_header(line: str): + def _has_imat_header(line: str) -> bool: HEADERS = [ "TIME STAMP,IMAGE TYPE,IMAGE COUNTER,COUNTS BM3 before image,COUNTS BM3 after image", "TIME STAMP IMAGE TYPE IMAGE COUNTER COUNTS BM3 before image COUNTS BM3 after image", @@ -108,7 +108,7 @@ def _has_imat_header(line: str): return line.strip() in HEADERS @classmethod - def _has_imat_data_line(cls, line: str): + def _has_imat_data_line(cls, line: str) -> bool: try: _ = cls.read_imat_date(line[:24]) except ValueError: diff --git a/mantidimaging/core/io/saver.py b/mantidimaging/core/io/saver.py index c7f790d52df..f28933f8d09 100644 --- a/mantidimaging/core/io/saver.py +++ b/mantidimaging/core/io/saver.py @@ -36,17 +36,20 @@ package_version = CheckVersion().get_version() -def write_fits(data: np.ndarray, filename: str, overwrite: bool = False, description: str | None = ""): +def write_fits(data: np.ndarray, filename: str, overwrite: bool = False, description: str | None = "") -> None: hdu = fits.PrimaryHDU(data) hdulist = fits.HDUList([hdu]) hdulist.writeto(filename, overwrite=overwrite) -def write_img(data: np.ndarray, filename: str, overwrite: bool = False, description: str | None = ""): +def write_img(data: np.ndarray, filename: str, overwrite: bool = False, description: str | None = "") -> None: tifffile.imwrite(filename, data, description=description, metadata=None, software="Mantid Imaging") -def write_nxs(data: np.ndarray, filename: str, projection_angles: np.ndarray | None = None, overwrite: bool = False): +def write_nxs(data: np.ndarray, + filename: str, + projection_angles: np.ndarray | None = None, + overwrite: bool = False) -> None: import h5py nxs = h5py.File(filename, 'w') @@ -177,7 +180,7 @@ def image_save(images: ImageStack, return names -def nexus_save(dataset: StrictDataset, path: str, sample_name: str, save_as_float: bool): +def nexus_save(dataset: StrictDataset, path: str, sample_name: str, save_as_float: bool) -> None: """ Uses information from a StrictDataset to create a NeXus file. :param dataset: The dataset to save as a NeXus file. @@ -199,7 +202,7 @@ def nexus_save(dataset: StrictDataset, path: str, sample_name: str, save_as_floa nexus_file.close() -def _nexus_save(nexus_file: h5py.File, dataset: StrictDataset, sample_name: str, save_as_float: bool): +def _nexus_save(nexus_file: h5py.File, dataset: StrictDataset, sample_name: str, save_as_float: bool) -> None: """ Takes a NeXus file and writes the StrictDataset information to it. :param nexus_file: The NeXus file. @@ -252,7 +255,7 @@ def _nexus_save(nexus_file: h5py.File, dataset: StrictDataset, sample_name: str, def _save_processed_data_to_nexus(nexus_file: h5py.File, dataset: StrictDataset, rotation_angle: h5py.Dataset, - image_key: h5py.Dataset, save_as_float: bool): + image_key: h5py.Dataset, save_as_float: bool) -> None: data = nexus_file.create_group(NEXUS_PROCESSED_DATA_PATH) data["rotation_angle"] = rotation_angle data["image_key"] = image_key @@ -266,7 +269,7 @@ def _save_processed_data_to_nexus(nexus_file: h5py.File, dataset: StrictDataset, process.create_dataset("version", data=np.bytes_(package_version)) -def _save_image_stacks_to_nexus(dataset: StrictDataset, data_group: h5py.Group, save_as_float: bool): +def _save_image_stacks_to_nexus(dataset: StrictDataset, data_group: h5py.Group, save_as_float: bool) -> None: combined_data_shape = (sum([len(arr) for arr in dataset.nexus_arrays]), ) + dataset.nexus_arrays[0].shape[1:] index = 0 @@ -305,7 +308,7 @@ def scale_row(row): return converted, factors -def _save_recon_to_nexus(nexus_file: h5py.File, recon: ImageStack, sample_path: str): +def _save_recon_to_nexus(nexus_file: h5py.File, recon: ImageStack, sample_path: str) -> None: """ Saves a recon to a NeXus file. :param nexus_file: The NeXus file. @@ -369,7 +372,7 @@ def _create_pixel_size_arrays(recon: ImageStack) -> tuple[np.ndarray, np.ndarray return x_arr, y_arr, z_arr -def _set_nx_class(group: h5py.Group, class_name: str): +def _set_nx_class(group: h5py.Group, class_name: str) -> None: """ Sets the NX_class attribute of data in a NeXus file. :param group: The h5py group. diff --git a/mantidimaging/core/net/help_pages.py b/mantidimaging/core/net/help_pages.py index b003b3e7dbf..0bbb5f3532f 100644 --- a/mantidimaging/core/net/help_pages.py +++ b/mantidimaging/core/net/help_pages.py @@ -9,13 +9,13 @@ SECTION_USER_GUIDE = f"{DOCS_BASE}/user_guide/" -def open_user_operation_docs(operation_name: str): +def open_user_operation_docs(operation_name: str) -> None: page_url = "operations/index" section = operation_name.lower().replace(" ", "-") open_help_webpage(SECTION_USER_GUIDE, page_url, section) -def open_help_webpage(section_url: str, page_url: str, section: str | None = None): +def open_help_webpage(section_url: str, page_url: str, section: str | None = None) -> None: if section is not None: url = f"{section_url}{page_url}.html#{section}" else: diff --git a/mantidimaging/core/parallel/shared.py b/mantidimaging/core/parallel/shared.py index ffa151def96..e5e98ad52a2 100644 --- a/mantidimaging/core/parallel/shared.py +++ b/mantidimaging/core/parallel/shared.py @@ -33,7 +33,7 @@ def run_compute_func(func: ComputeFuncType, num_operations: int, arrays: list[pu.SharedArray] | pu.SharedArray, params: dict[str, Any], - progress=None): + progress=None) -> None: if isinstance(arrays, pu.SharedArray): arrays = [arrays] all_data_in_shared_memory, data = _check_shared_mem_and_get_data(arrays) diff --git a/mantidimaging/core/parallel/utility.py b/mantidimaging/core/parallel/utility.py index 2e76516760f..2ed47e35c30 100644 --- a/mantidimaging/core/parallel/utility.py +++ b/mantidimaging/core/parallel/utility.py @@ -23,7 +23,7 @@ LOG = getLogger(__name__) -def enough_memory(shape, dtype): +def enough_memory(shape, dtype) -> bool: return full_size_KB(shape=shape, dtype=dtype) < system_free_memory().kb() @@ -103,7 +103,7 @@ def multiprocessing_necessary(shape: int, is_shared_data: bool) -> bool: return True -def execute_impl(img_num: int, partial_func: partial, is_shared_data: bool, progress: Progress, msg: str): +def execute_impl(img_num: int, partial_func: partial, is_shared_data: bool, progress: Progress, msg: str) -> None: task_name = f"{msg}" progress = Progress.ensure_instance(progress, num_steps=img_num, task_name=task_name) indices_list = range(img_num) @@ -128,7 +128,7 @@ def run_compute_func_impl(worker_func: Callable[[int], None], num_operations: int, is_shared_data: bool, progress=None, - msg: str = ""): + msg: str = "") -> None: task_name = f"{msg}" progress = Progress.ensure_instance(progress, num_steps=num_operations, task_name=task_name) indices_list = range(num_operations) diff --git a/mantidimaging/core/reconstruct/astra_recon.py b/mantidimaging/core/reconstruct/astra_recon.py index 3c031ea99e2..f049e722e01 100644 --- a/mantidimaging/core/reconstruct/astra_recon.py +++ b/mantidimaging/core/reconstruct/astra_recon.py @@ -24,11 +24,11 @@ # Full credit for following code to Daniil Kazantzev # Source: # https://github.com/dkazanc/ToMoBAR/blob/5990aaa264e2f08bd9b0069c8847e5021fbf2ee2/src/Python/tomobar/supp/astraOP.py#L20-L70 -def rotation_matrix2d(theta: float): +def rotation_matrix2d(theta: float) -> np.ndarray: return np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) -def vec_geom_init2d(angles_rad: ProjectionAngles, detector_spacing_x: float, center_rot_offset: float): +def vec_geom_init2d(angles_rad: ProjectionAngles, detector_spacing_x: float, center_rot_offset: float) -> np.ndarray: angles_value = angles_rad.value s0 = [0.0, -1.0] # source u0 = [detector_spacing_x, 0.0] # detector coordinates diff --git a/mantidimaging/core/reconstruct/base_recon.py b/mantidimaging/core/reconstruct/base_recon.py index 2805674f3c2..536a9031993 100644 --- a/mantidimaging/core/reconstruct/base_recon.py +++ b/mantidimaging/core/reconstruct/base_recon.py @@ -20,7 +20,7 @@ def find_cor(images: ImageStack, slice_idx: int, start_cor: float, recon_params: raise NotImplementedError("Base class call") @staticmethod - def prepare_sinogram(data: np.ndarray, recon_params: ReconstructionParameters): + def prepare_sinogram(data: np.ndarray, recon_params: ReconstructionParameters) -> np.ndarray: logged_data = BaseRecon.negative_log(data) if recon_params.beam_hardening_coefs is not None: coefs = np.array([0.0, 1.0] + recon_params.beam_hardening_coefs) diff --git a/mantidimaging/core/reconstruct/tomopy_recon.py b/mantidimaging/core/reconstruct/tomopy_recon.py index cfef8ff80dd..ef168f1a8f8 100644 --- a/mantidimaging/core/reconstruct/tomopy_recon.py +++ b/mantidimaging/core/reconstruct/tomopy_recon.py @@ -36,7 +36,7 @@ def single_sino(sino: np.ndarray, cor: ScalarCoR, proj_angles: ProjectionAngles, recon_params: ReconstructionParameters, - progress: Progress | None = None): + progress: Progress | None = None) -> np.ndarray: sino = BaseRecon.prepare_sinogram(sino, recon_params) volume = tomopy.recon(tomo=[sino], sinogram_order=True, @@ -51,7 +51,7 @@ def single_sino(sino: np.ndarray, def full(images: ImageStack, cors: list[ScalarCoR], recon_params: ReconstructionParameters, - progress: Progress | None = None): + progress: Progress | None = None) -> np.ndarray: """ Performs a volume reconstruction using sample data provided as sinograms. diff --git a/mantidimaging/core/rotation/data_model.py b/mantidimaging/core/rotation/data_model.py index e5745c3d095..d82cf7daeea 100644 --- a/mantidimaging/core/rotation/data_model.py +++ b/mantidimaging/core/rotation/data_model.py @@ -29,17 +29,17 @@ def __init__(self): self._cached_gradient = None self._cached_cor = None - def populate_slice_indices(self, begin, end, count, cor=0.0): + def populate_slice_indices(self, begin: int, end: int, count: int, cor: float = 0.0) -> None: self.clear_results() self._points = [Point(int(idx), cor) for idx in np.linspace(begin, end, count, dtype=int)] LOG.debug(f'Populated slice indices: {self.slices}') - def linear_regression(self): + def linear_regression(self) -> None: LOG.debug(f'Running linear regression with {self.num_points} points') self._cached_gradient, self._cached_cor, *_ = sp.stats.linregress(self.slices, self.cors) - def add_point(self, idx=None, slice_idx=0, cor=0.0): + def add_point(self, idx=None, slice_idx=0, cor=0.0) -> None: self.clear_results() if idx is None: @@ -47,7 +47,7 @@ def add_point(self, idx=None, slice_idx=0, cor=0.0): else: self._points.insert(idx, Point(slice_idx, cor)) - def set_point(self, idx, slice_idx: int | None = None, cor: float | None = None, reset_results=True): + def set_point(self, idx, slice_idx: int | None = None, cor: float | None = None, reset_results=True) -> None: if reset_results: self.clear_results() @@ -64,29 +64,29 @@ def _get_data_idx_from_slice_idx(self, slice_idx) -> int: return i raise ValueError(f"Slice {slice_idx} that is not in COR table") - def set_cor_at_slice(self, slice_idx: int, cor: float): + def set_cor_at_slice(self, slice_idx: int, cor: float) -> None: data_idx = self._get_data_idx_from_slice_idx(slice_idx) self.set_point(data_idx, cor=cor) - def remove_point(self, idx): + def remove_point(self, idx) -> None: self.clear_results() del self._points[idx] - def clear_points(self): + def clear_points(self) -> None: self._points = [] self.clear_results() - def clear_results(self): + def clear_results(self) -> None: self._cached_gradient = None self._cached_cor = None - def point(self, idx): + def point(self, idx: int) -> Point | None: return self._points[idx] if idx < self.num_points else None - def sort_points(self): + def sort_points(self) -> None: self._points.sort(key=lambda p: p.slice_index) - def get_cor_from_regression(self, slice_idx) -> float: + def get_cor_from_regression(self, slice_idx: int) -> float: cor = (self.gradient.value * slice_idx) + self.cor.value return cor @@ -104,11 +104,11 @@ def get_all_cors_from_regression(self, image_height) -> list[ScalarCoR]: return cors @property - def slices(self): + def slices(self) -> list[int]: return [p.slice_index for p in self._points] @property - def cors(self): + def cors(self) -> list[float]: return [float(p.cor) for p in self._points] @property @@ -130,19 +130,19 @@ def angle_in_degrees(self) -> Degrees: return Degrees(-np.rad2deg(np.arctan(self.gradient.value))) @property - def has_results(self): + def has_results(self) -> bool: return self._cached_gradient is not None and self._cached_cor is not None @property - def empty(self): + def empty(self) -> bool: return not self._points @property - def num_points(self): + def num_points(self) -> int: return len(self._points) @property - def stack_properties(self): + def stack_properties(self) -> dict: return { # TODO remove float casts const.COR_TILT_ROTATION_CENTRE: float(self.cor.value), @@ -152,7 +152,7 @@ def stack_properties(self): const.COR_TILT_ROTATION_CENTRES: self.cors } - def set_precalculated(self, cor: ScalarCoR, tilt: Degrees): + def set_precalculated(self, cor: ScalarCoR, tilt: Degrees) -> None: self._cached_cor = cor.value # reverse the tilt calculation to get the slope of the regression back self._cached_gradient = -np.tan(np.deg2rad(tilt.value)) diff --git a/mantidimaging/core/rotation/polyfit_correlation.py b/mantidimaging/core/rotation/polyfit_correlation.py index 5e06c6fdc31..290c63e0fe6 100644 --- a/mantidimaging/core/rotation/polyfit_correlation.py +++ b/mantidimaging/core/rotation/polyfit_correlation.py @@ -17,7 +17,7 @@ def do_calculate_correlation_err(store: np.ndarray, search_index: int, p0_and_180: tuple[np.ndarray, np.ndarray], - image_width: int): + image_width: int) -> None: """ Calculates squared sum error in the difference between the projection at 0 degrees, and the one at 180 degrees """ @@ -68,7 +68,7 @@ def find_center(images: ImageStack, progress: Progress) -> tuple[ScalarCoR, Degr return ScalarCoR(images.h_middle + -offset), theta -def compute_correlation_error(index: int, arrays: list[Any], params: dict[str, Any]): +def compute_correlation_error(index: int, arrays: list[Any], params: dict[str, Any]) -> None: min_correlation_error = arrays[0] shared_projections = arrays[1] shared_search_range = arrays[2] @@ -79,7 +79,7 @@ def compute_correlation_error(index: int, arrays: list[Any], params: dict[str, A (shared_projections[0], shared_projections[1]), image_width) -def _find_shift(images: ImageStack, search_range: range, min_correlation_error: np.ndarray, shift: np.ndarray): +def _find_shift(images: ImageStack, search_range: range, min_correlation_error: np.ndarray, shift: np.ndarray) -> None: # Then we just find the index of the minimum one (minimum error) min_correlation_error = np.transpose(min_correlation_error) # argmin returns a list of where the minimum argument is found