def total_dist_from_point(data: CachingDataStructure, start_point: tuple, max_prop_level: int = INF) -> tuple: """ For some CachingDataType, it calculates the distance in linear space for all different propagation levels. Returns a tuple of three arrays; (1) an array of propagation level, (2) summed linear distance and (3) summed linear distance divided by the number of pixels/voxels at that propagation level """ assert data.dim == len(start_point), \ "Dimensions of data and start point don't match" assert data.dim in [2, 3], "Data should be 2- or 3-dimensional" start_internal_index = data.internal_index(*start_point) get_indices = indices_for_prop_level_2d \ if data.dim == 2 \ else indices_for_prop_level_3d prop_level_end = propagation_level_end(start_point, data.shape) prop_levels = range(min(prop_level_end, max_prop_level)) ys1, ys2 = [], [] summ, num_of_voxels = 0, 0 for prop_level in prop_levels: for key in get_indices(prop_level, start_point): if data.valid_index(*key): summ += abs(data.internal_index(*key) - start_internal_index) num_of_voxels += 1 ys1.append(summ) ys2.append(summ / num_of_voxels) return np.array(prop_levels), np.array(ys1), np.array(ys2)
def convolution_3d(picture: CachingDataStructure, kernel: np.ndarray, cval: float = 0.0) -> CachingDataStructure: """Performs convolution for a 3D image. Pads borders with cval""" assert picture.dim == 3 ret_data = picture.empty_of_same_shape() kernel_dim = kernel.shape[0] pad = kernel_dim // 2 for (x, y, z) in picture.iter_keys(): ret_data[x, y, z] = convolve_voxel(picture, kernel, kernel_dim, pad, x, y, z, cval) return ret_data
def fmm(speed_func_arr: CachingDataStructure, start_point: tuple, end_point: tuple = None) -> CachingDataStructure: """ Computes the FMM for a n-dimensional speed function and a given start point If an end point is supplied, only the necessary coordinates will be computed. Otherwise all coordinates """ assert speed_func_arr.dim == len(start_point), \ f"Difference in dimensionality of speed func ({speed_func_arr.dim}) \ and start_point {start_point}" if speed_func_arr.dim == 2: return fmm_2d(speed_func_arr, start_point, end_point) elif speed_func_arr.dim == 3: return fmm_3d(speed_func_arr, start_point, end_point) # else keep doing general times = speed_func_arr.empty_of_same_shape() times.fill(INF) times[start_point] = 0.0 status = times.empty_of_same_shape() status.fill(False) status[start_point] = KNOWN heap_pointers = status.empty_of_same_shape() heap_pointers.fill(Neighbor(0, 0, False), dtype='object') neighbors = [] heapify(neighbors) first_run = True current_key = start_point while (current_key != end_point) and (first_run or neighbors): first_run = False # Get new neighbors coords = get_neighbors_general(status, current_key) # Calculate arrival times of newly added neighbors for key in coords: val_time = calc_time_general(speed_func_arr, times, key) if val_time < times[key]: if heap_pointers[key].valid: heap_pointers[key].valid = False times[key] = val_time heap_elm = Neighbor(val_time, key, True) heap_pointers[key] = heap_elm heappush(neighbors, heap_elm) # Find the smallet value of neighbours while neighbors: element = heappop(neighbors) if element.valid: break current_key = element.coords # Add coordinate of smallest value to known status[current_key] = KNOWN return times
def safe_get_times_3d(times: CachingDataStructure, i: int, j: int, k: int) -> float: """ Returns the saved time value in 3D array if is a valid index, otherwise infinity """ return INF if not times.valid_index(i, j, k) else times[i, j, k]
def get_neighbors_2d(status: CachingDataStructure, i: int, j: int) -> list: """ Returns a list of tuples of all coordinates that are direct neighbors, meaning the index is valid and they are not KNOWN """ coords = [] for (x, y) in [(i-1, j), (i+1, j), (i, j-1), (i, j+1)]: if status.valid_index(x, y) and not status[x, y]: # Not known coords.append((x, y)) return coords
def get_neighbors_general(status: CachingDataStructure, key: tuple) -> list: """ Returns a list of tuples of all coordinates that are direct neighbors, meaning the index is valid and they are not KNOWN """ coords = [] for key in get_direct_neighbour_coords_general(key): if status.valid_index(*key) and not status[key]: # Not known coords.append(key) return coords
def fmm_2d(speed_func_arr: CachingDataStructure, start_point: tuple, end_point: tuple = None)\ -> Tuple[CachingDataStructure, CachingDataStructure]: """ Computes the FMM for a 2D speed function and a given start point. If an end point is supplied, only the necessary coordinates will be computed. Otherwise all coordinates. Returns the time array and an array of neighbors that can be used to find the path to a given end point """ i, j = start_point times = speed_func_arr.empty_of_same_shape() times.fill(INF) times[i, j] = 0.0 # Status[i, j] is false if (i, j) is far or neighbour, true if (i, i) is # known status = times.empty_of_same_shape() status.fill(False, dtype=bool) status[i, j] = KNOWN heap_pointers = status.empty_of_same_shape() heap_pointers.fill(Neighbor(0, 0, False), dtype='object') neighbors = [] heapify(neighbors) first_run = True while (i, j) != end_point and (first_run or neighbors): first_run = False # Get new neighbors coords = get_neighbors_2d(status, i, j) # Calculate arrival times of newly added neighbors for (x, y) in coords: val_time, prev = calc_time_2d(speed_func_arr, times, x, y) if val_time < times[x, y]: if heap_pointers[x, y].valid: heap_pointers[x, y].valid = False times[x, y] = val_time heap_elm = Neighbor(val_time, [x, y], True, prev) heap_pointers[x, y] = heap_elm heappush(neighbors, heap_elm) # Find the smallet value of neighbours while neighbors: element = heappop(neighbors) if element.valid: break i, j = element.coords # Add coordinate of smallest value to known status[i, j] = KNOWN return times, heap_pointers
def bit_rev_copy(arr: CachingDataStructure, axis: int) -> CachingDataStructure: """ Copies an CachingDataStructure, but swaps the values along a given axis according to the reversed bit value of that index """ rev_arr = arr.empty_of_same_shape().fill(complex(0, 0), dtype='complex') n = arr.shape[0] bitlen = int(np.log2(n)) for key in rev_arr.iter_keys(): new_key = tuple_gen(rev_bit_str(key[axis], bitlen), key, axis) rev_arr[new_key] = arr[key] return rev_arr
def fmm_3d(speed_func_arr: CachingDataStructure, start_point: tuple, end_point: tuple = None) -> CachingDataStructure: """ Computes the FMM for a 3D speed function and a given start point. If an end point is supplied, only the necessary coordinates will be computed. Otherwise all coordinates """ i, j, k = start_point times = speed_func_arr.empty_of_same_shape() times.fill(INF) times[i, j, k] = 0.0 status = times.empty_of_same_shape() status.fill(False) status[i, j, k] = KNOWN heap_pointers = status.empty_of_same_shape() heap_pointers.fill(Neighbor(0, 0, False), dtype='object') neighbors = [] heapify(neighbors) first_run = True while (i, j, k) != end_point and (first_run or neighbors): first_run = False # Get new neighbors coords = get_neighbors_3d(status, i, j, k) # Calculate arrival times of newly added neighbors for (x, y, z) in coords: val_time = calc_time_3d(speed_func_arr, times, x, y, z) if val_time < times[x, y, z]: if heap_pointers[x, y, z].valid: heap_pointers[x, y, z].valid = False times[x, y, z] = val_time heap_elm = Neighbor(val_time, [x, y, z], True) heap_pointers[x, y, z] = heap_elm heappush(neighbors, heap_elm) # Find the smallet value of neighbours while neighbors: element = heappop(neighbors) if element.valid: break i, j, k = element.coords # Add coordinate of smallest value to known status[i, j, k] = KNOWN return times
def convolve_pixel(picture: CachingDataStructure, kernel: np.ndarray, kernel_dim: int, pad: int, x: int, y: int, cval: float) -> CachingDataStructure: """Performs spatial convolution at a single pixel in a 2D image""" pixel = 0.0 for i in range(kernel_dim): for j in range(kernel_dim): x_val, y_val = x + i - pad, y + j - pad if picture.valid_index(x_val, y_val): pixel += kernel[i, j] * picture[x_val, y_val] else: pixel += kernel[i, j] * cval return pixel
def convolve_voxel(picture: CachingDataStructure, kernel: np.ndarray, kernel_dim: int, pad: int, x: int, y: int, z: int, cval: float) -> float: """Performs spatial convolution at a single voxel in a 3D image""" voxel = 0.0 for i in range(kernel_dim): for j in range(kernel_dim): for h in range(kernel_dim): x_val, y_val, z_val = x + i - pad, y + j - pad, z + h - pad if picture.valid_index(x_val, y_val, z_val): voxel += kernel[i, j, h] * picture[x_val, y_val, z_val] else: voxel += kernel[i, j, h] * cval return voxel
def matmul_rec(A: CachingDataStructure, B: CachingDataStructure, r: int) -> CachingDataStructure: """ Entry point for calling the recursive matrix multiplication to some recursion depth r. A and B should both be n x n CachingDataStructure where n = 2^t. Returns the result of the matrix multiplication """ assert A.dim == 2 and B.dim == 2, \ f"Arrays should be 2D but got dimensions {A.dim} and {B.dim}" assert A.shape == B.shape, \ f"Arrays should be same shape but got {A.shape} and {B.shape}" C = B.empty_of_same_shape() n = B.shape[0] mm_recursive(A, B, C, (0, 0), (0, 0), (0, 0), n, r) return C
def bit_rev_inplace(arr: CachingDataStructure, axis: int) \ -> CachingDataStructure: """ Swaps the values along a given axis according to the reversed bit value of that index """ N = arr.shape[0] bitlen = int(np.log2(N)) for key in arr.iter_keys(): rev_bit_idx = rev_bit_str(key[axis], bitlen) if key[axis] >= rev_bit_idx: # Make sure to only swap elements once continue new_key = tuple_gen(rev_bit_idx, key, axis) arr[key], arr[new_key] = arr[new_key], arr[key] return arr
def safe_get_times_general(times: CachingDataStructure, key: tuple) -> float: """ Returns the saved time value in an n-dimensional array if is a valid index, otherwise infinity """ return INF if not times.valid_index(*key) else times[key]