Coverage for src/arraybridge/slice_processing.py: 98%

1""" 

2Shared slice-by-slice processing logic for all memory types. 

3 

4This module provides a single implementation of slice-by-slice processing 

5that works for all memory types, eliminating duplication across dtype wrappers. 

6""" 

7 

8from arraybridge.converters import detect_memory_type 

9from arraybridge.stack_utils import stack_slices, unstack_slices 

10from arraybridge.utils import _get_device_id 

11 

12 

13def process_slices(image, func, args, kwargs, gpu_id=None): 

14 """ 

15 Process a 3D array slice-by-slice using the provided function. 

16 

17 This function handles: 

18 - Unstacking 3D arrays into 2D slices 

19 - Processing each slice independently 

20 - Handling functions that return tuples (main output + special outputs) 

21 - Stacking results back into 3D arrays in the returned main-output framework 

22 - Combining special outputs from all slices 

23 

24 Args: 

25 image: 3D array to process 

26 func: Function to apply to each slice 

27 args: Positional arguments to pass to func 

28 kwargs: Keyword arguments to pass to func 

29 gpu_id: Optional GPU device ID override. If not provided, attempts 

30 to derive from the input image and falls back to 0. 

31 

32 Returns: 

33 Processed 3D array, or tuple of (processed_3d_array, special_outputs...) 

34 if func returns tuples 

35 """ 

36 # Detect memory type and use proper OpenHCS utilities 

37 memory_type = detect_memory_type(image) 

38 if gpu_id is None: 

39 detected_gpu_id = _get_device_id(image, memory_type) 

40 gpu_id = 0 if detected_gpu_id is None else detected_gpu_id 

41 

42 # Unstack 3D array into 2D slices 

43 slices_2d = unstack_slices(image, memory_type, gpu_id) 

44 

45 # Process each slice and handle special outputs 

46 main_outputs = [] 

47 special_outputs_list = [] 

48 returns_tuple = None 

49 tuple_arity = None 

50 

51 for slice_index, slice_2d in enumerate(slices_2d): 

52 slice_result = func(slice_2d, *args, **kwargs) 

53 

54 # Check if result is a tuple (indicating special outputs) 

55 result_is_tuple = isinstance(slice_result, tuple) 

56 if returns_tuple is None: 

57 returns_tuple = result_is_tuple 

58 elif result_is_tuple != returns_tuple: 

59 raise TypeError( 

60 "Slice processing cannot mix tuple and non-tuple results; " 

61 f"slice {slice_index} returned {type(slice_result).__name__}." 

62 ) 

63 

64 if result_is_tuple: 

65 if not slice_result: 

66 raise ValueError("Slice processing result tuples cannot be empty") 

67 if tuple_arity is None: 

68 tuple_arity = len(slice_result) 

69 elif len(slice_result) != tuple_arity: 

70 raise ValueError( 

71 "Slice processing requires every result tuple to have the " 

72 f"same arity; slice {slice_index} returned {len(slice_result)}, " 

73 f"expected {tuple_arity}." 

74 ) 

75 main_outputs.append(slice_result[0]) # First element is main output 

76 special_outputs_list.append(slice_result[1:]) # Rest are special outputs 

77 else: 

78 main_outputs.append(slice_result) # Single output 

79 

80 # Stack main outputs in the framework returned by the callable. The input 

81 # framework owns unstacking only; decorators may declare a different output 

82 # framework for the per-slice function. 

83 if not main_outputs: 

84 raise ValueError("Slice processing produced no main outputs to stack") 

85 output_memory_type = detect_memory_type(main_outputs[0]) 

86 output_gpu_id = _get_device_id(main_outputs[0], output_memory_type) 

87 result = stack_slices( 

88 main_outputs, 

89 output_memory_type, 

90 0 if output_gpu_id is None else output_gpu_id, 

91 ) 

92 

93 # If we have special outputs, combine them and return tuple 

94 if special_outputs_list: 

95 # Combine special outputs from all slices 

96 combined_special_outputs = [] 

97 num_special_outputs = len(special_outputs_list[0]) 

98 

99 for i in range(num_special_outputs): 

100 # Collect the i-th special output from all slices 

101 special_output_values = [slice_outputs[i] for slice_outputs in special_outputs_list] 

102 combined_special_outputs.append(special_output_values) 

103 

104 # Return tuple: (stacked_main_output, combined_special_output1, # noqa: E501 

105 # combined_special_output2, ...) 

106 return (result, *combined_special_outputs) 

107 

108 return result