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]