Coverage for src/arraybridge/stack_utils.py: 81%

1""" 

2Stack utilities module for OpenHCS. 

3 

4This module provides functions for stacking 2D slices into a 3D array 

5and unstacking a 3D array into 2D slices, with explicit memory type handling. 

6 

7This module enforces Clause 278 — Mandatory 3D Output Enforcement: 

8All functions must return a 3D array of shape [Z, Y, X], even when operating 

9on a single 2D slice. No logic may check, coerce, or infer rank at unstack time. 

10""" 

11 

12import logging 

13from typing import Any 

14 

15from arraybridge.converters import detect_memory_type 

16from arraybridge.types import MemoryType 

17 

18logger = logging.getLogger(__name__) 

19 

20# 🔍 MEMORY CONVERSION LOGGING: Test log to verify logger is working 

21logger.debug("🔄 STACK_UTILS: Module loaded - memory conversion logging enabled") 

22 

23 

24def _is_2d(data: Any) -> bool: 

25 """ 

26 Check if data is a 2D array. 

27 

28 Args: 

29 data: Data to check 

30 

31 Returns: 

32 True if data is 2D, False otherwise 

33 """ 

34 # Check if data has a shape attribute 

35 if not hasattr(data, "shape"): 

36 return False 

37 

38 # Check if shape has length 2 

39 return len(data.shape) == 2 

40 

41 

42def _is_3d(data: Any) -> bool: 

43 """ 

44 Check if data is a 3D array. 

45 

46 Args: 

47 data: Data to check 

48 

49 Returns: 

50 True if data is 3D, False otherwise 

51 """ 

52 # Check if data has a shape attribute 

53 if not hasattr(data, "shape"): 

54 return False 

55 

56 # Check if shape has length 3 

57 return len(data.shape) == 3 

58 

59 

60def _enforce_gpu_device_requirements(memory_type: str, gpu_id: int) -> None: 

61 """ 

62 Enforce GPU device requirements. 

63 

64 Args: 

65 memory_type: The memory type 

66 gpu_id: The GPU device ID 

67 

68 Raises: 

69 ValueError: If gpu_id is negative 

70 """ 

71 mem_type = MemoryType(memory_type) 

72 if mem_type.is_gpu: 

73 mem_type.require_device(gpu_id) 

74 

75 

76def stack_slices(slices: list[Any], memory_type: str, gpu_id: int) -> Any: 

77 """ 

78 Stack 2D slices into a 3D array with the specified memory type. 

79 

80 STRICT VALIDATION: Assumes all slices are 2D arrays. 

81 No automatic handling of improper inputs. 

82 

83 Args: 

84 slices: List of 2D slices (numpy arrays, cupy arrays, torch tensors, etc.) 

85 memory_type: The memory type to use for the stacked array (REQUIRED) 

86 gpu_id: The target GPU device ID (REQUIRED) 

87 

88 Returns: 

89 A 3D array with the specified memory type of shape [Z, Y, X] 

90 

91 Raises: 

92 ValueError: If memory_type is not supported or slices is empty 

93 ValueError: If gpu_id is negative for GPU memory types 

94 ValueError: If slices are not 2D arrays 

95 MemoryConversionError: If conversion fails 

96 """ 

97 if not slices: 

98 raise ValueError("Cannot stack empty list of slices") 

99 

100 # Verify all slices are 2D 

101 for i, slice_data in enumerate(slices): 

102 if not _is_2d(slice_data): 

103 raise ValueError(f"Slice at index {i} is not a 2D array. All slices must be 2D.") 

104 

105 # Check GPU requirements 

106 _enforce_gpu_device_requirements(memory_type, gpu_id) 

107 

108 # Convert each slice and enforce the requested framework-local device. 

109 conversion_count = 0 

110 mem_type = MemoryType(memory_type) 

111 converted_slices = [] 

112 for slice_data in slices: 

113 source_type = detect_memory_type(slice_data) 

114 if source_type == memory_type: 

115 converted_data = mem_type.move_to_device(slice_data, gpu_id) 

116 else: 

117 from arraybridge.converters import convert_memory 

118 

119 converted_data = convert_memory( 

120 data=slice_data, 

121 source_type=source_type, 

122 target_type=memory_type, 

123 gpu_id=gpu_id, 

124 ) 

125 conversion_count += 1 

126 converted_slices.append(converted_data) 

127 

128 result = mem_type.stack_arrays(converted_slices, gpu_id) 

129 

130 # 🔍 MEMORY CONVERSION LOGGING: Only log when conversions happen or issues occur 

131 if conversion_count > 0: 

132 logger.debug( 

133 f"🔄 STACK_SLICES: Converted {conversion_count}/{len(slices)} " 

134 f"slices to {memory_type}" 

135 ) 

136 # Silent success for no-conversion cases to reduce log pollution 

137 

138 return result 

139 

140 

141def unstack_slices( 

142 array: Any, memory_type: str, gpu_id: int, validate_slices: bool = True 

143) -> list[Any]: 

144 """ 

145 Split a 3D array into 2D slices along axis 0 and convert to the specified memory type. 

146 

147 STRICT VALIDATION: Input must be a 3D array. No automatic handling of improper inputs. 

148 

149 Args: 

150 array: 3D array to split - MUST BE 3D 

151 memory_type: The memory type to use for the output slices (REQUIRED) 

152 gpu_id: The target GPU device ID (REQUIRED) 

153 validate_slices: If True, validates that each extracted slice is 2D 

154 

155 Returns: 

156 List of 2D slices in the specified memory type 

157 

158 Raises: 

159 ValueError: If array is not 3D 

160 ValueError: If validate_slices is True and any extracted slice is not 2D 

161 ValueError: If gpu_id is negative for GPU memory types 

162 ValueError: If memory_type is not supported 

163 MemoryConversionError: If conversion fails 

164 """ 

165 # Detect input type and check if conversion is needed 

166 input_type = detect_memory_type(array) 

167 getattr(array, "shape", "unknown") 

168 

169 # Verify the array is 3D - fail loudly if not 

170 if not _is_3d(array): 

171 raise ValueError(f"Array must be 3D, got shape {getattr(array, 'shape', 'unknown')}") 

172 

173 # Check GPU requirements 

174 _enforce_gpu_device_requirements(memory_type, gpu_id) 

175 

176 # Convert to target memory type 

177 source_type = input_type # Reuse already detected type 

178 

179 # Direct conversion 

180 if source_type == memory_type: 

181 array = MemoryType(memory_type).move_to_device(array, gpu_id) 

182 else: 

183 # Convert and log the conversion 

184 from arraybridge.converters import convert_memory 

185 

186 logger.debug(f"🔄 UNSTACK_SLICES: Converting array - {source_type} → {memory_type}") 

187 array = convert_memory( 

188 data=array, source_type=source_type, target_type=memory_type, gpu_id=gpu_id 

189 ) 

190 

191 # Extract slices along axis 0 (already in the target memory type) 

192 slices = [array[i] for i in range(array.shape[0])] 

193 

194 # Validate that all extracted slices are 2D if requested 

195 if validate_slices: 

196 for i, slice_data in enumerate(slices): 

197 if not _is_2d(slice_data): 

198 raise ValueError( 

199 f"Extracted slice at index {i} is not 2D. " 

200 f"This indicates a malformed 3D array." 

201 ) 

202 

203 # 🔍 MEMORY CONVERSION LOGGING: Only log conversions or issues 

204 if source_type != memory_type: 

205 logger.debug(f"🔄 UNSTACK_SLICES: Converted and extracted {len(slices)} slices") 

206 elif len(slices) == 0: 

207 logger.warning("🔄 UNSTACK_SLICES: No slices extracted (empty array)") 

208 # Silent success for no-conversion cases to reduce log pollution 

209 

210 return slices